From f82f940738e4338fb75c2d54260a4c38c6c4eb1a Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Thu, 9 Apr 2026 16:13:50 -0400 Subject: [PATCH 1/3] Converted window type and cleanup policies to enums --- .../rs/asap_types/src/aggregation_config.rs | 18 ++--- .../rs/asap_types/src/capability_matching.rs | 9 +-- .../dependencies/rs/asap_types/src/enums.rs | 58 ++++++++++++++++ .../rs/asap_types/src/inference_config.rs | 11 ++-- asap-planner-rs/src/output/generator.rs | 18 ++--- asap-planner-rs/src/output/sql_generator.rs | 8 +-- asap-planner-rs/src/planner/logics.rs | 66 ++++++++++++++----- asap-planner-rs/src/planner/single_query.rs | 12 ++-- .../src/planner/sql_single_query.rs | 4 +- .../benches/simple_store_bench.rs | 10 +-- .../src/bin/bench_precompute_sketch.rs | 3 +- .../src/bin/test_e2e_precompute.rs | 7 +- asap-query-engine/src/data_model/enums.rs | 2 +- .../src/engines/simple_engine.rs | 11 ++-- .../precompute_engine/accumulator_factory.rs | 5 +- .../src/precompute_engine/worker.rs | 9 +-- .../src/tests/capability_matching_tests.rs | 25 ++++--- .../plan_execution_temporal_tests.rs | 31 ++++----- .../src/tests/query_equivalence_tests.rs | 10 ++- .../src/tests/sql_pattern_matching_tests.rs | 4 +- .../src/tests/store_correctness_tests.rs | 11 ++-- .../tests/test_utilities/config_builders.rs | 14 ++-- .../tests/test_utilities/engine_factories.rs | 19 +++--- .../tests/e2e_precompute_equivalence.rs | 7 +- 24 files changed, 236 insertions(+), 136 deletions(-) diff --git a/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs b/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs index b709eaa..6820d4b 100644 --- a/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs @@ -3,7 +3,7 @@ use serde_json::Value; use serde_yaml; use std::collections::HashMap; -use crate::enums::QueryLanguage; +use crate::enums::{QueryLanguage, WindowType}; use crate::traits::SerializableToSink; use crate::utils::normalize_spatial_filter; use promql_utilities::data_model::KeyByLabelNames; @@ -19,9 +19,9 @@ pub struct AggregationConfig { pub rollup_labels: KeyByLabelNames, pub original_yaml: String, - pub window_size: u64, // Window size in seconds (e.g., 900s for 15m) - pub slide_interval: u64, // Slide/hop interval in seconds (e.g., 30s) - pub window_type: String, // "tumbling" or "sliding" + pub window_size: u64, // Window size in seconds (e.g., 900s for 15m) + pub slide_interval: u64, // Slide/hop interval in seconds (e.g., 30s) + pub window_type: WindowType, // Tumbling or Sliding pub spatial_filter: String, pub spatial_filter_normalized: String, @@ -60,7 +60,7 @@ impl AggregationConfig { original_yaml: String, window_size: u64, slide_interval: u64, - window_type: String, + window_type: WindowType, spatial_filter: String, metric: String, num_aggregates_to_retain: Option, @@ -148,7 +148,8 @@ impl AggregationConfig { .get("windowType") .and_then(|v| v.as_str()) .unwrap_or("tumbling") - .to_string(); + .parse::() + .unwrap_or_default(); let slide_interval = data .get("slideInterval") @@ -270,7 +271,8 @@ impl AggregationConfig { .get("windowType") .and_then(|v| v.as_str()) .unwrap_or("tumbling") - .to_string(); + .parse::() + .unwrap_or_default(); let slide_interval = aggregation_data .get("slideInterval") @@ -348,7 +350,7 @@ impl SerializableToSink for AggregationConfig { "originalYaml": self.original_yaml, "windowSize": self.window_size, "slideInterval": self.slide_interval, - "windowType": self.window_type, + "windowType": self.window_type.to_string(), "spatialFilter": self.spatial_filter, "metric": self.metric, }); diff --git a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs index 1dce691..2dbb1a2 100644 --- a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs +++ b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs @@ -6,6 +6,7 @@ use promql_utilities::query_logics::enums::Statistic; use tracing::{debug, warn}; use crate::aggregation_config::{AggregationConfig, AggregationIdInfo}; +use crate::enums::WindowType; use crate::query_requirements::QueryRequirements; use crate::utils::normalize_spatial_filter; @@ -76,9 +77,9 @@ pub fn window_compatible(config: &AggregationConfig, data_range_ms: Option) if window_ms == 0 || range == 0 { return false; } - match config.window_type.as_str() { - "sliding" => range == window_ms, - _ => range % window_ms == 0, // tumbling (or unknown — treat as tumbling) + match config.window_type { + WindowType::Sliding => range == window_ms, + WindowType::Tumbling => range % window_ms == 0, } } @@ -288,7 +289,7 @@ mod tests { original_yaml: String::new(), window_size: window_size_s, slide_interval: window_size_s, - window_type: window_type.to_string(), + window_type: window_type.parse::().unwrap_or_default(), spatial_filter: spatial_filter.to_string(), spatial_filter_normalized, metric: metric.to_string(), diff --git a/asap-common/dependencies/rs/asap_types/src/enums.rs b/asap-common/dependencies/rs/asap_types/src/enums.rs index 3b5bb6a..c20e88b 100644 --- a/asap-common/dependencies/rs/asap_types/src/enums.rs +++ b/asap-common/dependencies/rs/asap_types/src/enums.rs @@ -1,3 +1,6 @@ +use std::fmt; +use std::str::FromStr; + #[derive(clap::ValueEnum, Clone, Copy, Debug, PartialEq)] #[allow(non_camel_case_types)] pub enum QueryLanguage { @@ -23,3 +26,58 @@ pub enum CleanupPolicy { /// Never clean up aggregates NoCleanup, } + +impl fmt::Display for CleanupPolicy { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + CleanupPolicy::CircularBuffer => write!(f, "circular_buffer"), + CleanupPolicy::ReadBased => write!(f, "read_based"), + CleanupPolicy::NoCleanup => write!(f, "no_cleanup"), + } + } +} + +impl FromStr for CleanupPolicy { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + "circular_buffer" => Ok(CleanupPolicy::CircularBuffer), + "read_based" => Ok(CleanupPolicy::ReadBased), + "no_cleanup" => Ok(CleanupPolicy::NoCleanup), + _ => Err(format!("Unknown cleanup policy: '{s}'")), + } + } +} + +/// Window type for streaming aggregations. +#[derive( + Clone, Debug, Copy, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize, +)] +#[serde(rename_all = "snake_case")] +pub enum WindowType { + #[default] + Tumbling, + Sliding, +} + +impl fmt::Display for WindowType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + WindowType::Tumbling => write!(f, "tumbling"), + WindowType::Sliding => write!(f, "sliding"), + } + } +} + +impl FromStr for WindowType { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "tumbling" => Ok(WindowType::Tumbling), + "sliding" => Ok(WindowType::Sliding), + _ => Err(format!("Unknown window type: '{s}'")), + } + } +} diff --git a/asap-common/dependencies/rs/asap_types/src/inference_config.rs b/asap-common/dependencies/rs/asap_types/src/inference_config.rs index 5778a6c..8ddb3ca 100644 --- a/asap-common/dependencies/rs/asap_types/src/inference_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/inference_config.rs @@ -172,15 +172,12 @@ impl InferenceConfig { ) })?; - match name { - "circular_buffer" => Ok(CleanupPolicy::CircularBuffer), - "read_based" => Ok(CleanupPolicy::ReadBased), - "no_cleanup" => Ok(CleanupPolicy::NoCleanup), - _ => Err(anyhow::anyhow!( + name.parse::().map_err(|_| { + anyhow::anyhow!( "Invalid cleanup policy: '{}'. Valid options: circular_buffer, read_based, no_cleanup", name - )), - } + ) + }) } fn parse_query_configs( diff --git a/asap-planner-rs/src/output/generator.rs b/asap-planner-rs/src/output/generator.rs index ec2c9df..69b168f 100644 --- a/asap-planner-rs/src/output/generator.rs +++ b/asap-planner-rs/src/output/generator.rs @@ -29,7 +29,9 @@ pub fn generate_plan( .as_ref() .and_then(|c| c.policy.as_deref()) .unwrap_or("read_based"); - let cleanup_policy = parse_cleanup_policy(cleanup_policy_str)?; + let cleanup_policy = cleanup_policy_str.parse::().map_err(|_| { + ControllerError::PlannerError(format!("Unknown cleanup policy: {}", cleanup_policy_str)) + })?; // Validate no duplicate queries let mut seen_queries = std::collections::HashSet::new(); @@ -187,18 +189,6 @@ fn collect_binary_leaf_entries( Ok(Some(all_entries)) } -pub fn parse_cleanup_policy(s: &str) -> Result { - match s { - "circular_buffer" => Ok(CleanupPolicy::CircularBuffer), - "read_based" => Ok(CleanupPolicy::ReadBased), - "no_cleanup" => Ok(CleanupPolicy::NoCleanup), - other => Err(ControllerError::PlannerError(format!( - "Unknown cleanup policy: {}", - other - ))), - } -} - pub fn key_by_labels_to_yaml(labels: &KeyByLabelNames) -> YamlValue { YamlValue::Sequence( labels @@ -278,7 +268,7 @@ pub fn build_aggregation_entry(id: u32, cfg: &IntermediateAggConfig) -> YamlValu ); map.insert( YamlValue::String("windowType".to_string()), - YamlValue::String(cfg.window_type.clone()), + YamlValue::String(cfg.window_type.to_string()), ); YamlValue::Mapping(map) diff --git a/asap-planner-rs/src/output/sql_generator.rs b/asap-planner-rs/src/output/sql_generator.rs index a3802ef..73851f0 100644 --- a/asap-planner-rs/src/output/sql_generator.rs +++ b/asap-planner-rs/src/output/sql_generator.rs @@ -6,9 +6,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use crate::config::input::SQLControllerConfig; use crate::error::ControllerError; -use crate::output::generator::{ - build_aggregation_entry, build_queries_yaml, parse_cleanup_policy, GeneratorOutput, -}; +use crate::output::generator::{build_aggregation_entry, build_queries_yaml, GeneratorOutput}; use crate::planner::single_query::IntermediateAggConfig; use crate::planner::sql_single_query::SQLSingleQueryProcessor; use crate::StreamingEngine; @@ -35,7 +33,9 @@ pub fn generate_sql_plan( .as_ref() .and_then(|c| c.policy.as_deref()) .unwrap_or("read_based"); - let cleanup_policy = parse_cleanup_policy(cleanup_policy_str)?; + let cleanup_policy = cleanup_policy_str.parse::().map_err(|_| { + ControllerError::PlannerError(format!("Unknown cleanup policy: {}", cleanup_policy_str)) + })?; // Validate T % data_ingestion_interval == 0 for qg in &config.query_groups { diff --git a/asap-planner-rs/src/planner/logics.rs b/asap-planner-rs/src/planner/logics.rs index 6e8160f..10dc0e9 100644 --- a/asap-planner-rs/src/planner/logics.rs +++ b/asap-planner-rs/src/planner/logics.rs @@ -1,5 +1,5 @@ use crate::config::input::SketchParameterOverrides; -use asap_types::enums::CleanupPolicy; +use asap_types::enums::{CleanupPolicy, WindowType}; use promql_utilities::ast_matching::PromQLMatchResult; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::{QueryPatternType, Statistic}; @@ -60,12 +60,12 @@ fn set_tumbling_window_parameters( QueryPatternType::OnlyTemporal | QueryPatternType::OneTemporalOneSpatial => { config.window_size = effective_repeat; config.slide_interval = effective_repeat; - config.window_type = "tumbling".to_string(); + config.window_type = WindowType::Tumbling; } QueryPatternType::OnlySpatial => { config.window_size = prometheus_scrape_interval; config.slide_interval = prometheus_scrape_interval; - config.window_type = "tumbling".to_string(); + config.window_type = WindowType::Tumbling; } } } @@ -75,7 +75,7 @@ fn set_tumbling_window_parameters( pub struct IntermediateWindowConfig { pub window_size: u64, pub slide_interval: u64, - pub window_type: String, + pub window_type: WindowType, } /// Shared sketch parameter builder used by both PromQL and SQL paths. @@ -213,7 +213,7 @@ pub fn get_cleanup_param( query_pattern_type: QueryPatternType, match_result: &PromQLMatchResult, t_repeat: u64, - window_type: &str, + window_type: WindowType, range_duration: u64, step: u64, ) -> Result { @@ -236,7 +236,7 @@ pub fn get_cleanup_param( .ok_or_else(|| "No range_vector token found".to_string())? }; - if window_type == "sliding" { + if window_type == WindowType::Sliding { let result = if is_range_query { range_duration / step + 1 } else { @@ -358,7 +358,7 @@ mod tests { pt, &mr, 300, - "tumbling", + WindowType::Tumbling, 0, 0, ) @@ -376,7 +376,7 @@ mod tests { pt, &mr, 300, - "tumbling", + WindowType::Tumbling, 3600, 30, ) @@ -388,8 +388,16 @@ mod tests { fn cleanup_param_read_based_spatial_instant_query() { let (pt, mr) = match_query("sum(some_metric)"); // lookback_buckets = ceil(300/300) = 1, num_steps = 1 → result = 1 - let result = - get_cleanup_param(CleanupPolicy::ReadBased, pt, &mr, 300, "tumbling", 0, 0).unwrap(); + let result = get_cleanup_param( + CleanupPolicy::ReadBased, + pt, + &mr, + 300, + WindowType::Tumbling, + 0, + 0, + ) + .unwrap(); assert_eq!(result, 1); } @@ -398,9 +406,16 @@ mod tests { let (pt, mr) = match_query("sum(some_metric)"); // lookback_buckets = ceil(300/30) = 10, num_steps = 3600/30 + 1 = 121 // result = 10 * 121 = 1210 - let result = - get_cleanup_param(CleanupPolicy::ReadBased, pt, &mr, 300, "tumbling", 3600, 30) - .unwrap(); + let result = get_cleanup_param( + CleanupPolicy::ReadBased, + pt, + &mr, + 300, + WindowType::Tumbling, + 3600, + 30, + ) + .unwrap(); assert_eq!(result, 1210); } @@ -410,16 +425,31 @@ mod tests { assert_eq!(pt, QueryPatternType::OnlyTemporal); // t_lookback = 5m = 300s (from [5m] range vector), range_duration=0, step=0 // effective_repeat = 60, ceil((300 + 0) / 60) = 5 - let result = - get_cleanup_param(CleanupPolicy::CircularBuffer, pt, &mr, 60, "tumbling", 0, 0) - .unwrap(); + let result = get_cleanup_param( + CleanupPolicy::CircularBuffer, + pt, + &mr, + 60, + WindowType::Tumbling, + 0, + 0, + ) + .unwrap(); assert_eq!(result, 5); } #[test] fn cleanup_param_no_cleanup_returns_error() { let (pt, mr) = match_query("sum(some_metric)"); - let result = get_cleanup_param(CleanupPolicy::NoCleanup, pt, &mr, 300, "tumbling", 0, 0); + let result = get_cleanup_param( + CleanupPolicy::NoCleanup, + pt, + &mr, + 300, + WindowType::Tumbling, + 0, + 0, + ); assert!(result.is_err()); } @@ -432,7 +462,7 @@ mod tests { pt, &mr, 300, - "tumbling", + WindowType::Tumbling, 3600, 0, ); diff --git a/asap-planner-rs/src/planner/single_query.rs b/asap-planner-rs/src/planner/single_query.rs index 55906ae..221a7cb 100644 --- a/asap-planner-rs/src/planner/single_query.rs +++ b/asap-planner-rs/src/planner/single_query.rs @@ -1,4 +1,4 @@ -use asap_types::enums::CleanupPolicy; +use asap_types::enums::{CleanupPolicy, WindowType}; use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; use promql_utilities::data_model::KeyByLabelNames; @@ -57,7 +57,7 @@ fn strip_parens(expr: &promql_parser::parser::Expr) -> &promql_parser::parser::E pub struct IntermediateAggConfig { pub aggregation_type: String, pub aggregation_sub_type: String, - pub window_type: String, + pub window_type: WindowType, pub window_size: u64, pub slide_interval: u64, pub spatial_filter: String, @@ -331,7 +331,7 @@ impl SingleQueryProcessor { pattern_type, &match_result, self.t_repeat, - &window_cfg.window_type, + window_cfg.window_type, self.range_duration, self.step, ) @@ -424,7 +424,7 @@ pub fn build_agg_configs_for_statistics( configs.push(IntermediateAggConfig { aggregation_type: "DeltaSetAggregator".to_string(), aggregation_sub_type: String::new(), - window_type: window_cfg.window_type.clone(), + window_type: window_cfg.window_type, window_size: window_cfg.window_size, slide_interval: window_cfg.slide_interval, spatial_filter: spatial_filter.to_string(), @@ -442,7 +442,7 @@ pub fn build_agg_configs_for_statistics( configs.push(IntermediateAggConfig { aggregation_type: agg_type, aggregation_sub_type: agg_sub_type, - window_type: window_cfg.window_type.clone(), + window_type: window_cfg.window_type, window_size: window_cfg.window_size, slide_interval: window_cfg.slide_interval, spatial_filter: spatial_filter.to_string(), @@ -469,7 +469,7 @@ mod tests { IntermediateAggConfig { aggregation_type: "MultipleIncrease".to_string(), aggregation_sub_type: "rate".to_string(), - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, window_size: 300, slide_interval: 300, spatial_filter: String::new(), diff --git a/asap-planner-rs/src/planner/sql_single_query.rs b/asap-planner-rs/src/planner/sql_single_query.rs index 2963f57..53edbdd 100644 --- a/asap-planner-rs/src/planner/sql_single_query.rs +++ b/asap-planner-rs/src/planner/sql_single_query.rs @@ -1,6 +1,6 @@ use std::collections::HashSet; -use asap_types::enums::CleanupPolicy; +use asap_types::enums::{CleanupPolicy, WindowType}; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::{QueryTreatmentType, Statistic}; use sql_utilities::ast_matching::sqlhelper::Table; @@ -198,7 +198,7 @@ fn compute_sql_window( IntermediateWindowConfig { window_size, slide_interval: window_size, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, } } diff --git a/asap-query-engine/benches/simple_store_bench.rs b/asap-query-engine/benches/simple_store_bench.rs index 403d96f..27a907c 100644 --- a/asap-query-engine/benches/simple_store_bench.rs +++ b/asap-query-engine/benches/simple_store_bench.rs @@ -27,7 +27,7 @@ use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; use promql_utilities::data_model::KeyByLabelNames; use query_engine_rust::data_model::{ - AggregateCore, CleanupPolicy, KeyByLabelValues, LockStrategy, StreamingConfig, + AggregateCore, CleanupPolicy, KeyByLabelValues, LockStrategy, StreamingConfig, WindowType, }; use query_engine_rust::precompute_operators::{DatasketchesKLLAccumulator, SumAccumulator}; use query_engine_rust::stores::simple_map_store::legacy::{ @@ -134,10 +134,10 @@ fn make_agg_config( KeyByLabelNames::empty(), KeyByLabelNames::empty(), "".to_string(), - 60, // window_size (seconds) - 60, // slide_interval (seconds) - "tumbling".to_string(), // window_type - "".to_string(), // spatial_filter + 60, // window_size (seconds) + 60, // slide_interval (seconds) + WindowType::Tumbling, // window_type + "".to_string(), // spatial_filter metric.to_string(), num_aggregates_to_retain, read_count_threshold, diff --git a/asap-query-engine/src/bin/bench_precompute_sketch.rs b/asap-query-engine/src/bin/bench_precompute_sketch.rs index 2d98ac8..f99cfb4 100644 --- a/asap-query-engine/src/bin/bench_precompute_sketch.rs +++ b/asap-query-engine/src/bin/bench_precompute_sketch.rs @@ -1,4 +1,5 @@ use asap_types::aggregation_config::AggregationConfig; +use asap_types::enums::WindowType; use clap::Parser; use prost::Message; use query_engine_rust::data_model::{ @@ -135,7 +136,7 @@ fn make_kll_streaming_config( String::new(), window_size_secs, window_size_secs, - "tumbling".to_string(), + WindowType::Tumbling, "bench_metric".to_string(), "bench_metric".to_string(), None, diff --git a/asap-query-engine/src/bin/test_e2e_precompute.rs b/asap-query-engine/src/bin/test_e2e_precompute.rs index a09357e..3ef70fd 100644 --- a/asap-query-engine/src/bin/test_e2e_precompute.rs +++ b/asap-query-engine/src/bin/test_e2e_precompute.rs @@ -9,6 +9,7 @@ //! cargo run --bin test_e2e_precompute use asap_types::aggregation_config::AggregationConfig; +use asap_types::enums::WindowType; use prost::Message; use query_engine_rust::data_model::{LockStrategy, QueryLanguage, StreamingConfig}; use query_engine_rust::drivers::ingest::prometheus_remote_write::{ @@ -584,9 +585,9 @@ fn make_sum_agg_config( slide_interval_secs: u64, ) -> AggregationConfig { let window_type = if slide_interval_secs == 0 || slide_interval_secs == window_size_secs { - "tumbling" + WindowType::Tumbling } else { - "sliding" + WindowType::Sliding }; AggregationConfig::new( agg_id, @@ -599,7 +600,7 @@ fn make_sum_agg_config( String::new(), window_size_secs, slide_interval_secs, - window_type.to_string(), + window_type, "bench_metric".to_string(), "bench_metric".to_string(), None, diff --git a/asap-query-engine/src/data_model/enums.rs b/asap-query-engine/src/data_model/enums.rs index ad5a090..3cc571f 100644 --- a/asap-query-engine/src/data_model/enums.rs +++ b/asap-query-engine/src/data_model/enums.rs @@ -9,7 +9,7 @@ pub enum StreamingEngine { Arroyo, } -pub use asap_types::enums::{CleanupPolicy, QueryLanguage}; +pub use asap_types::enums::{CleanupPolicy, QueryLanguage, WindowType}; #[derive(clap::ValueEnum, Clone, Debug, PartialEq)] pub enum QueryProtocol { diff --git a/asap-query-engine/src/engines/simple_engine.rs b/asap-query-engine/src/engines/simple_engine.rs index 5f52b94..58f693c 100644 --- a/asap-query-engine/src/engines/simple_engine.rs +++ b/asap-query-engine/src/engines/simple_engine.rs @@ -17,6 +17,7 @@ use tracing::{debug, warn}; use crate::AggregateCore; +use asap_types::enums::WindowType; use asap_types::query_requirements::QueryRequirements; use asap_types::utils::normalize_spatial_filter; use promql_utilities::ast_matching::{PromQLMatchResult, PromQLPattern, PromQLPatternBuilder}; @@ -592,8 +593,8 @@ impl SimpleEngine { ) })?; - let window_type = &aggregation_config_for_value.window_type; - let is_exact_query = window_type == "sliding"; + let window_type = aggregation_config_for_value.window_type; + let is_exact_query = window_type == WindowType::Sliding; // Determine start/end for values query based on window type let (values_start, values_end) = if is_exact_query { @@ -720,9 +721,9 @@ impl SimpleEngine { let merge_start_time = Instant::now(); let window_type = if plan.values_query.is_exact_query { - "sliding" + WindowType::Sliding } else { - "tumbling" + WindowType::Tumbling }; let merged_values = if plan.values_query.is_exact_query { @@ -755,7 +756,7 @@ impl SimpleEngine { let merge_duration = merge_start_time.elapsed(); debug!( "[LATENCY] Precomputed output processing ({}): {:.2}ms, resulted in {} merged outputs", - if window_type == "sliding" { + if window_type == WindowType::Sliding { "no merge" } else { "merge" diff --git a/asap-query-engine/src/precompute_engine/accumulator_factory.rs b/asap-query-engine/src/precompute_engine/accumulator_factory.rs index c1b6c9c..77d2ca7 100644 --- a/asap-query-engine/src/precompute_engine/accumulator_factory.rs +++ b/asap-query-engine/src/precompute_engine/accumulator_factory.rs @@ -649,6 +649,7 @@ pub fn create_accumulator_updater(config: &AggregationConfig) -> Box, ) -> AggregationConfig { let window_type = if slide_secs == 0 || slide_secs == window_secs { - "tumbling" + WindowType::Tumbling } else { - "sliding" + WindowType::Sliding }; AggregationConfig::new( id, @@ -646,7 +647,7 @@ mod tests { String::new(), window_secs, slide_secs, - window_type.to_string(), + window_type, metric.to_string(), metric.to_string(), None, @@ -1300,7 +1301,7 @@ aggregations: String::new(), 60, 0, - "tumbling".to_string(), + WindowType::Tumbling, "http_requests_total".to_string(), "http_requests_total".to_string(), Some(60), diff --git a/asap-query-engine/src/tests/capability_matching_tests.rs b/asap-query-engine/src/tests/capability_matching_tests.rs index 7dbc75c..9bec8ad 100644 --- a/asap-query-engine/src/tests/capability_matching_tests.rs +++ b/asap-query-engine/src/tests/capability_matching_tests.rs @@ -6,7 +6,7 @@ use crate::data_model::{ AggregationConfig, AggregationReference, CleanupPolicy, InferenceConfig, PrecomputedOutput, - PromQLSchema, QueryConfig, QueryLanguage, SchemaConfig, StreamingConfig, + PromQLSchema, QueryConfig, QueryLanguage, SchemaConfig, StreamingConfig, WindowType, }; use crate::engines::simple_engine::SimpleEngine; use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; @@ -27,7 +27,7 @@ fn make_agg_config( metric: &str, agg_type: &str, window_size_s: u64, - window_type: &str, + window_type: WindowType, grouping: &[&str], ) -> AggregationConfig { AggregationConfig { @@ -41,7 +41,7 @@ fn make_agg_config( original_yaml: String::new(), window_size: window_size_s, slide_interval: window_size_s, - window_type: window_type.to_string(), + window_type, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: metric.to_string(), @@ -161,7 +161,7 @@ fn engine_with_query_config( /// capability matching should route to it and return a valid context. #[test] fn capability_fallback_fires_when_no_config() { - let agg = make_agg_config(1, "cpu", "Sum", 300, "tumbling", &[]); + let agg = make_agg_config(1, "cpu", "Sum", 300, WindowType::Tumbling, &[]); let engine = engine_no_query_configs("cpu", &[], vec![agg]); // sum_over_time(cpu[5m]) — 5 min = 300 s matches the 300 s tumbling config @@ -178,7 +178,7 @@ fn capability_fallback_fires_when_no_config() { /// We verify by giving the config a different agg_id than any compatible-by-type config. #[test] fn config_path_takes_priority_over_capability_matching() { - let agg = make_agg_config(42, "cpu", "Sum", 300, "tumbling", &[]); + let agg = make_agg_config(42, "cpu", "Sum", 300, WindowType::Tumbling, &[]); let engine = engine_with_query_config("cpu", &[], agg, "sum_over_time(cpu[5m])"); let ctx = engine @@ -193,7 +193,14 @@ fn config_path_takes_priority_over_capability_matching() { /// KLL aggregation when no query_configs are present. #[test] fn quantile_different_values_resolve_to_same_aggregation() { - let kll = make_agg_config(7, "latency", "DatasketchesKLL", 300, "tumbling", &[]); + let kll = make_agg_config( + 7, + "latency", + "DatasketchesKLL", + 300, + WindowType::Tumbling, + &[], + ); let engine = engine_no_query_configs("latency", &[], vec![kll]); let q50 = engine.build_query_execution_context_promql( @@ -224,7 +231,7 @@ fn quantile_different_values_resolve_to_same_aggregation() { #[test] fn no_match_returns_none() { // KLL config present, but query asks for Sum — incompatible - let kll = make_agg_config(1, "cpu", "DatasketchesKLL", 300, "tumbling", &[]); + let kll = make_agg_config(1, "cpu", "DatasketchesKLL", 300, WindowType::Tumbling, &[]); let engine = engine_no_query_configs("cpu", &[], vec![kll]); let ctx = @@ -238,8 +245,8 @@ fn no_match_returns_none() { /// When multiple compatible aggregations exist, the largest window should be preferred. #[test] fn priority_largest_window_wins() { - let small = make_agg_config(1, "cpu", "Sum", 300, "tumbling", &[]); - let large = make_agg_config(2, "cpu", "Sum", 900, "tumbling", &[]); + let small = make_agg_config(1, "cpu", "Sum", 300, WindowType::Tumbling, &[]); + let large = make_agg_config(2, "cpu", "Sum", 900, WindowType::Tumbling, &[]); let engine = engine_no_query_configs("cpu", &[], vec![small, large]); // sum_over_time(cpu[15m]) = 900 s — both 300 s and 900 s configs match (900 = 3×300), diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs index 609d747..5b0d5e1 100644 --- a/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs @@ -10,6 +10,7 @@ use std::collections::HashMap; #[cfg(test)] mod tests { use super::*; + use crate::data_model::WindowType; use crate::precompute_operators::DatasketchesKLLAccumulator; use crate::tests::test_utilities::engine_factories::*; @@ -51,7 +52,7 @@ mod tests { data, query, 5, - "tumbling", + WindowType::Tumbling, ); let results = execute_new_plan(&engine, query, QUERY_TIME).await; @@ -81,7 +82,7 @@ mod tests { data, query, 5, - "tumbling", + WindowType::Tumbling, ); let results = execute_new_plan(&engine, query, QUERY_TIME).await; @@ -118,7 +119,7 @@ mod tests { data, query, 5, - "tumbling", + WindowType::Tumbling, ); let results = execute_new_plan(&engine, query, QUERY_TIME).await; @@ -155,7 +156,7 @@ mod tests { data, query, 5, - "tumbling", + WindowType::Tumbling, ); let results = execute_new_plan(&engine, query, QUERY_TIME).await; @@ -193,7 +194,7 @@ mod tests { data, query, 5, - "tumbling", + WindowType::Tumbling, ); let results = execute_new_plan(&engine, query, QUERY_TIME).await; @@ -226,7 +227,7 @@ mod tests { vec![], query, 5, - "tumbling", + WindowType::Tumbling, ); let results = execute_new_plan(&engine, query, QUERY_TIME).await; @@ -251,7 +252,7 @@ mod tests { )], query, 5, - "tumbling", + WindowType::Tumbling, ); let context = engine @@ -275,7 +276,7 @@ mod tests { )], query, 5, - "tumbling", + WindowType::Tumbling, ); let context = engine @@ -324,7 +325,7 @@ mod tests { data, query, 5, - "tumbling", + WindowType::Tumbling, ); let results = execute_new_plan(&engine, query, QUERY_TIME).await; @@ -376,7 +377,7 @@ mod tests { data, query, 5, - "tumbling", + WindowType::Tumbling, ); let results = execute_new_plan(&engine, query, QUERY_TIME).await; @@ -415,7 +416,7 @@ mod tests { )], query, 5, - "tumbling", + WindowType::Tumbling, ); let context = engine @@ -443,7 +444,7 @@ mod tests { )], query, 5, - "tumbling", + WindowType::Tumbling, ); let context = engine @@ -467,7 +468,7 @@ mod tests { vec![], query, 5, - "tumbling", + WindowType::Tumbling, ); let results = execute_new_plan(&engine, query, QUERY_TIME).await; @@ -503,7 +504,7 @@ mod tests { data, query, 5, - "tumbling", + WindowType::Tumbling, ); assert_old_new_match(&engine, query, QUERY_TIME).await; @@ -533,7 +534,7 @@ mod tests { data, query, 5, - "tumbling", + WindowType::Tumbling, ); assert_old_new_match(&engine, query, QUERY_TIME).await; diff --git a/asap-query-engine/src/tests/query_equivalence_tests.rs b/asap-query-engine/src/tests/query_equivalence_tests.rs index 80aa93e..d202787 100644 --- a/asap-query-engine/src/tests/query_equivalence_tests.rs +++ b/asap-query-engine/src/tests/query_equivalence_tests.rs @@ -7,7 +7,7 @@ //! timestamp calculation, and aggregation selection - WITHOUT actually executing //! queries against a store. -use crate::data_model::QueryLanguage; +use crate::data_model::{QueryLanguage, WindowType}; use crate::engines::simple_engine::SimpleEngine; use crate::stores::{Store, TimestampedBucketsMap}; use crate::tests::test_utilities::{assert_execution_context_equivalent, TestConfigBuilder}; @@ -90,7 +90,13 @@ mod tests { let (promql_config, sql_config, streaming_config) = TestConfigBuilder::new("cpu_usage") .with_grouping_labels(grouping_labels) .with_scrape_interval(scrape_interval) - .add_temporal_query(promql_query, sql_query, 1, window_seconds, "tumbling") + .add_temporal_query( + promql_query, + sql_query, + 1, + window_seconds, + WindowType::Tumbling, + ) .build_both(); // Create engines (they won't query the store) diff --git a/asap-query-engine/src/tests/sql_pattern_matching_tests.rs b/asap-query-engine/src/tests/sql_pattern_matching_tests.rs index b5bc8f4..e0e566c 100644 --- a/asap-query-engine/src/tests/sql_pattern_matching_tests.rs +++ b/asap-query-engine/src/tests/sql_pattern_matching_tests.rs @@ -7,7 +7,7 @@ mod tests { use crate::data_model::{ AggregationConfig, AggregationReference, CleanupPolicy, InferenceConfig, QueryConfig, - QueryLanguage, SchemaConfig, StreamingConfig, + QueryLanguage, SchemaConfig, StreamingConfig, WindowType, }; use crate::engines::simple_engine::SimpleEngine; use crate::stores::simple_map_store::SimpleMapStore; @@ -63,7 +63,7 @@ mod tests { original_yaml: String::new(), window_size: window_secs, slide_interval: window_secs, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: "cpu_usage".to_string(), diff --git a/asap-query-engine/src/tests/store_correctness_tests.rs b/asap-query-engine/src/tests/store_correctness_tests.rs index efc41bd..62f7215 100644 --- a/asap-query-engine/src/tests/store_correctness_tests.rs +++ b/asap-query-engine/src/tests/store_correctness_tests.rs @@ -28,7 +28,8 @@ //! | `contract_global` | `LockStrategy::Global` | use crate::data_model::{ - CleanupPolicy, KeyByLabelValues, LockStrategy, Measurement, SerializableToSink, StreamingConfig, + CleanupPolicy, KeyByLabelValues, LockStrategy, Measurement, SerializableToSink, + StreamingConfig, WindowType, }; use crate::precompute_operators::{ CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, DatasketchesKLLAccumulator, @@ -59,10 +60,10 @@ fn make_agg_config( KeyByLabelNames::empty(), KeyByLabelNames::empty(), "".to_string(), - 60, // window_size (seconds) - 60, // slide_interval (seconds) - "tumbling".to_string(), // window_type - "".to_string(), // spatial_filter + 60, // window_size (seconds) + 60, // slide_interval (seconds) + WindowType::Tumbling, // window_type + "".to_string(), // spatial_filter "cpu_usage".to_string(), num_aggregates_to_retain, read_count_threshold, diff --git a/asap-query-engine/src/tests/test_utilities/config_builders.rs b/asap-query-engine/src/tests/test_utilities/config_builders.rs index 72135dc..0183f8f 100644 --- a/asap-query-engine/src/tests/test_utilities/config_builders.rs +++ b/asap-query-engine/src/tests/test_utilities/config_builders.rs @@ -5,7 +5,7 @@ use crate::data_model::{ AggregationConfig, AggregationReference, CleanupPolicy, InferenceConfig, PromQLSchema, - QueryConfig, SchemaConfig, StreamingConfig, + QueryConfig, SchemaConfig, StreamingConfig, WindowType, }; use promql_utilities::data_model::KeyByLabelNames; use sql_utilities::sqlhelper::{SQLSchema, Table}; @@ -75,7 +75,7 @@ impl TestConfigBuilder { sql: &str, agg_id: u64, window_seconds: u64, - window_type: &str, // "sliding" or "tumbling" + window_type: WindowType, ) -> Self { // Add PromQL query config let promql_config = QueryConfig::new(promql.to_string()) @@ -99,7 +99,7 @@ impl TestConfigBuilder { original_yaml: String::new(), window_size: window_seconds, slide_interval: window_seconds, - window_type: window_type.to_string(), + window_type, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: self.metric.clone(), @@ -136,7 +136,7 @@ impl TestConfigBuilder { original_yaml: String::new(), window_size: self.scrape_interval, slide_interval: self.scrape_interval, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: self.metric.clone(), @@ -179,7 +179,7 @@ impl TestConfigBuilder { original_yaml: String::new(), window_size: window_seconds, slide_interval: window_seconds, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: self.metric.clone(), @@ -296,7 +296,7 @@ mod tests { "SELECT SUM(value) FROM cpu_usage WHERE time BETWEEN NOW() AND DATEADD(s, -10, NOW()) GROUP BY L1, L2, L3, L4", 1, 10, - "tumbling", + WindowType::Tumbling, ) .build(); @@ -317,7 +317,7 @@ mod tests { assert!(streaming_config.get_aggregation_config(1).is_some()); let agg_config = streaming_config.get_aggregation_config(1).unwrap(); assert_eq!(agg_config.window_size, 10); - assert_eq!(agg_config.window_type, "tumbling"); + assert_eq!(agg_config.window_type, WindowType::Tumbling); } #[test] diff --git a/asap-query-engine/src/tests/test_utilities/engine_factories.rs b/asap-query-engine/src/tests/test_utilities/engine_factories.rs index b1a27e1..51aaed2 100644 --- a/asap-query-engine/src/tests/test_utilities/engine_factories.rs +++ b/asap-query-engine/src/tests/test_utilities/engine_factories.rs @@ -8,6 +8,7 @@ use crate::data_model::{ AggregationConfig, AggregationReference, CleanupPolicy, InferenceConfig, KeyByLabelValues, PrecomputedOutput, PromQLSchema, QueryConfig, QueryLanguage, SchemaConfig, StreamingConfig, + WindowType, }; use crate::engines::query_result::InstantVectorElement; use crate::engines::simple_engine::SimpleEngine; @@ -81,7 +82,7 @@ pub fn create_engine_single_pop_with_aggregated( original_yaml: String::new(), window_size: 1, slide_interval: 1, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: metric.to_string(), @@ -178,7 +179,7 @@ pub fn create_engine_dual_input( original_yaml: String::new(), window_size: 1, slide_interval: 1, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: metric.to_string(), @@ -201,7 +202,7 @@ pub fn create_engine_dual_input( original_yaml: String::new(), window_size: 1, slide_interval: 1, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: metric.to_string(), @@ -294,7 +295,7 @@ pub fn create_engine_two_metrics( original_yaml: String::new(), window_size: 1, slide_interval: 1, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: metric_a.to_string(), @@ -316,7 +317,7 @@ pub fn create_engine_two_metrics( original_yaml: String::new(), window_size: 1, slide_interval: 1, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: metric_b.to_string(), @@ -419,7 +420,7 @@ pub fn create_engine_three_metrics( original_yaml: String::new(), window_size: 1, slide_interval: 1, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: metric.to_string(), @@ -500,7 +501,7 @@ pub fn create_engine_multi_timestamp( original_yaml: String::new(), window_size: 1, slide_interval: 1, - window_type: "tumbling".to_string(), + window_type: WindowType::Tumbling, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: metric.to_string(), @@ -563,7 +564,7 @@ pub fn create_engine_multi_timestamp_with_window( data: Vec<(u64, Option>, Box)>, promql_query: &str, window_size: u64, - window_type: &str, + window_type: WindowType, ) -> SimpleEngine { let grouping_label_strings: Vec = grouping_labels.iter().map(|s| s.to_string()).collect(); @@ -580,7 +581,7 @@ pub fn create_engine_multi_timestamp_with_window( original_yaml: String::new(), window_size, slide_interval: 1, - window_type: window_type.to_string(), + window_type, spatial_filter: String::new(), spatial_filter_normalized: String::new(), metric: metric.to_string(), diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index f193e53..53dd572 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -8,6 +8,7 @@ //! 4. Drains captured outputs and verifies equivalence with ArroYo-format accumulators use asap_types::aggregation_config::AggregationConfig; +use asap_types::enums::WindowType; use flate2::{write::GzEncoder, Compression}; use prost::Message; use serde_json::json; @@ -38,9 +39,9 @@ fn make_agg_config( grouping: Vec<&str>, ) -> AggregationConfig { let window_type = if slide_secs == 0 || slide_secs == window_secs { - "tumbling" + WindowType::Tumbling } else { - "sliding" + WindowType::Sliding }; AggregationConfig::new( id, @@ -55,7 +56,7 @@ fn make_agg_config( String::new(), window_secs, slide_secs, - window_type.to_string(), + window_type, metric.to_string(), metric.to_string(), None, From d5ae6528a23340da5942a1c7d4640ef702457cb1 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Thu, 9 Apr 2026 16:34:08 -0400 Subject: [PATCH 2/3] added more enums --- .../src/ast_matching/promql_pattern.rs | 16 ++- .../src/query_logics/enums.rs | 136 ++++++++++++++++++ .../src/query_logics/logics.rs | 50 +++++-- .../src/query_logics/parsing.rs | 10 +- asap-planner-rs/src/planner/patterns.rs | 72 ++++++---- asap-planner-rs/src/planner/single_query.rs | 32 +++-- .../src/engines/simple_engine.rs | 108 ++++++++------ 7 files changed, 325 insertions(+), 99 deletions(-) diff --git a/asap-common/dependencies/rs/promql_utilities/src/ast_matching/promql_pattern.rs b/asap-common/dependencies/rs/promql_utilities/src/ast_matching/promql_pattern.rs index 18874e4..a757f07 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/ast_matching/promql_pattern.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/ast_matching/promql_pattern.rs @@ -448,11 +448,11 @@ impl PromQLPattern { debug!("Collecting aggregation token as: {}", collect_as); let modifier = match &agg.modifier { Some(LabelModifier::Include(labels)) => Some(AggregationModifier { - modifier_type: "by".to_string(), + modifier_type: AggregationModifierType::By, labels: labels.labels.clone(), }), Some(LabelModifier::Exclude(labels)) => Some(AggregationModifier { - modifier_type: "without".to_string(), + modifier_type: AggregationModifierType::Without, labels: labels.labels.clone(), }), None => None, @@ -844,16 +844,24 @@ impl Default for PromQLMatchResult { } } +/// Whether a PromQL aggregation modifier is `by` or `without`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum AggregationModifierType { + By, + Without, +} + /// Represents aggregation modifiers like "by" or "without" #[derive(Debug, Clone, Serialize)] pub struct AggregationModifier { - pub modifier_type: String, // "by" or "without" + pub modifier_type: AggregationModifierType, pub labels: Vec, } impl AggregationModifier { /// Create a new AggregationModifier - pub fn new(modifier_type: String, labels: Vec) -> Self { + pub fn new(modifier_type: AggregationModifierType, labels: Vec) -> Self { Self { modifier_type, labels, diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs index 7369f79..883c3a6 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs @@ -111,6 +111,142 @@ impl std::fmt::Display for QueryResultType { } } +/// A PromQL function that produces a vector-valued result. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum PromQLFunction { + Rate, + Increase, + SumOverTime, + CountOverTime, + AvgOverTime, + MinOverTime, + MaxOverTime, + QuantileOverTime, +} + +impl PromQLFunction { + pub fn as_str(self) -> &'static str { + match self { + PromQLFunction::Rate => "rate", + PromQLFunction::Increase => "increase", + PromQLFunction::SumOverTime => "sum_over_time", + PromQLFunction::CountOverTime => "count_over_time", + PromQLFunction::AvgOverTime => "avg_over_time", + PromQLFunction::MinOverTime => "min_over_time", + PromQLFunction::MaxOverTime => "max_over_time", + PromQLFunction::QuantileOverTime => "quantile_over_time", + } + } + + /// Returns `true` for functions whose result requires approximate pre-aggregation. + pub fn is_approximate(self) -> bool { + matches!( + self, + PromQLFunction::QuantileOverTime + | PromQLFunction::SumOverTime + | PromQLFunction::CountOverTime + | PromQLFunction::AvgOverTime + ) + } +} + +impl std::fmt::Display for PromQLFunction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +impl std::str::FromStr for PromQLFunction { + type Err = String; + fn from_str(s: &str) -> Result { + match s { + "rate" => Ok(PromQLFunction::Rate), + "increase" => Ok(PromQLFunction::Increase), + "sum_over_time" => Ok(PromQLFunction::SumOverTime), + "count_over_time" => Ok(PromQLFunction::CountOverTime), + "avg_over_time" => Ok(PromQLFunction::AvgOverTime), + "min_over_time" => Ok(PromQLFunction::MinOverTime), + "max_over_time" => Ok(PromQLFunction::MaxOverTime), + "quantile_over_time" => Ok(PromQLFunction::QuantileOverTime), + other => Err(format!("Unknown PromQL function: '{other}'")), + } + } +} + +/// A PromQL aggregation operator. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum AggregationOperator { + Sum, + Count, + Avg, + Quantile, + Min, + Max, + Topk, +} + +impl AggregationOperator { + pub fn as_str(self) -> &'static str { + match self { + AggregationOperator::Sum => "sum", + AggregationOperator::Count => "count", + AggregationOperator::Avg => "avg", + AggregationOperator::Quantile => "quantile", + AggregationOperator::Min => "min", + AggregationOperator::Max => "max", + AggregationOperator::Topk => "topk", + } + } + + /// Returns the `Statistic` values required to answer this operator. + /// `Avg` needs both `Sum` and `Count`. + pub fn to_statistics(self) -> Vec { + match self { + AggregationOperator::Avg => vec![Statistic::Sum, Statistic::Count], + AggregationOperator::Sum => vec![Statistic::Sum], + AggregationOperator::Count => vec![Statistic::Count], + AggregationOperator::Quantile => vec![Statistic::Quantile], + AggregationOperator::Min => vec![Statistic::Min], + AggregationOperator::Max => vec![Statistic::Max], + AggregationOperator::Topk => vec![Statistic::Topk], + } + } + + /// Returns `true` for operators whose result requires approximate pre-aggregation. + pub fn is_approximate(self) -> bool { + matches!( + self, + AggregationOperator::Quantile + | AggregationOperator::Sum + | AggregationOperator::Count + | AggregationOperator::Avg + | AggregationOperator::Topk + ) + } +} + +impl std::fmt::Display for AggregationOperator { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +impl std::str::FromStr for AggregationOperator { + type Err = String; + fn from_str(s: &str) -> Result { + match s { + "sum" => Ok(AggregationOperator::Sum), + "count" => Ok(AggregationOperator::Count), + "avg" => Ok(AggregationOperator::Avg), + "quantile" => Ok(AggregationOperator::Quantile), + "min" => Ok(AggregationOperator::Min), + "max" => Ok(AggregationOperator::Max), + "topk" => Ok(AggregationOperator::Topk), + other => Err(format!("Unknown aggregation operator: '{other}'")), + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs index fd2093d..0684605 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs @@ -1,4 +1,6 @@ -use crate::query_logics::enums::{QueryTreatmentType, Statistic}; +use crate::query_logics::enums::{ + AggregationOperator, PromQLFunction, QueryTreatmentType, Statistic, +}; use tracing::debug; /// Map statistic to precompute operator based on treatment type @@ -80,20 +82,24 @@ pub fn does_precompute_operator_support_subpopulations( } } -/// Check if temporal and spatial aggregations are collapsible +/// Check if temporal and spatial aggregations are collapsible. /// Based on Python implementation in promql_utilities/query_logics/logics.py -pub fn get_is_collapsable(temporal_aggregation: &str, spatial_aggregation: &str) -> bool { +pub fn get_is_collapsable( + temporal_aggregation: PromQLFunction, + spatial_aggregation: AggregationOperator, +) -> bool { debug!( "Checking if temporal aggregation '{}' and spatial aggregation '{}' are collapsable", temporal_aggregation, spatial_aggregation ); match spatial_aggregation { - "sum" => matches!( + AggregationOperator::Sum => matches!( temporal_aggregation, - "sum_over_time" | "count_over_time" // Note: "increase" and "rate" are commented out in Python + // Note: Increase and Rate are commented out in the Python reference + PromQLFunction::SumOverTime | PromQLFunction::CountOverTime ), - "min" => temporal_aggregation == "min_over_time", - "max" => temporal_aggregation == "max_over_time", + AggregationOperator::Min => temporal_aggregation == PromQLFunction::MinOverTime, + AggregationOperator::Max => temporal_aggregation == PromQLFunction::MaxOverTime, _ => false, } } @@ -171,11 +177,29 @@ mod tests { #[test] fn test_get_is_collapsable() { - assert!(get_is_collapsable("sum_over_time", "sum")); - assert!(get_is_collapsable("count_over_time", "sum")); - assert!(get_is_collapsable("min_over_time", "min")); - assert!(get_is_collapsable("max_over_time", "max")); - assert!(!get_is_collapsable("min_over_time", "sum")); - assert!(!get_is_collapsable("unknown", "sum")); + assert!(get_is_collapsable( + PromQLFunction::SumOverTime, + AggregationOperator::Sum + )); + assert!(get_is_collapsable( + PromQLFunction::CountOverTime, + AggregationOperator::Sum + )); + assert!(get_is_collapsable( + PromQLFunction::MinOverTime, + AggregationOperator::Min + )); + assert!(get_is_collapsable( + PromQLFunction::MaxOverTime, + AggregationOperator::Max + )); + assert!(!get_is_collapsable( + PromQLFunction::MinOverTime, + AggregationOperator::Sum + )); + assert!(!get_is_collapsable( + PromQLFunction::Rate, + AggregationOperator::Sum + )); } } diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs index c37bb48..b57c699 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs @@ -3,6 +3,7 @@ use core::panic; use promql_parser::parser::Expr; 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::{QueryPatternType, Statistic}; @@ -115,21 +116,20 @@ pub fn get_spatial_aggregation_output_labels( }; debug!( - "Modifier type: {}, labels: {:?}", + "Modifier type: {:?}, labels: {:?}", modifier.modifier_type, modifier.labels ); - match modifier.modifier_type.as_str() { - "by" => { + match modifier.modifier_type { + AggregationModifierType::By => { debug!("Processing 'by' modifier"); // Return only the labels specified in "by" clause KeyByLabelNames::new(modifier.labels.clone()) } - "without" => { + AggregationModifierType::Without => { debug!("Processing 'without' modifier"); // Return all labels except those specified in "without" clause let without_labels = KeyByLabelNames::new(modifier.labels.clone()); all_labels.difference(&without_labels) } - _ => panic!("Invalid aggregation modifier"), } } diff --git a/asap-planner-rs/src/planner/patterns.rs b/asap-planner-rs/src/planner/patterns.rs index 31b8554..d2de5af 100644 --- a/asap-planner-rs/src/planner/patterns.rs +++ b/asap-planner-rs/src/planner/patterns.rs @@ -1,5 +1,7 @@ use promql_utilities::ast_matching::{PromQLPattern, PromQLPatternBuilder}; -use promql_utilities::query_logics::enums::QueryPatternType; +use promql_utilities::query_logics::enums::{ + AggregationOperator, PromQLFunction, QueryPatternType, +}; /// Build all 5 patterns in priority order: ONLY_TEMPORAL (2), ONLY_SPATIAL (1), ONE_TEMPORAL_ONE_SPATIAL (2) pub fn build_patterns() -> Vec<(QueryPatternType, PromQLPattern)> { @@ -7,9 +9,47 @@ pub fn build_patterns() -> Vec<(QueryPatternType, PromQLPattern)> { let range_vector_pattern = || PromQLPatternBuilder::matrix_selector(metric_pattern(), None, Some("range_vector")); + // Temporal functions that produce a single-value result (no quantile phi arg) + let temporal_funcs: Vec<&str> = [ + PromQLFunction::SumOverTime, + PromQLFunction::CountOverTime, + PromQLFunction::AvgOverTime, + PromQLFunction::MinOverTime, + PromQLFunction::MaxOverTime, + PromQLFunction::Increase, + PromQLFunction::Rate, + ] + .map(PromQLFunction::as_str) + .to_vec(); + + // Aggregation operators used in spatial and spatial-of-temporal patterns + let spatial_ops: Vec<&str> = [ + AggregationOperator::Sum, + AggregationOperator::Count, + AggregationOperator::Avg, + AggregationOperator::Quantile, + AggregationOperator::Min, + AggregationOperator::Max, + AggregationOperator::Topk, + ] + .map(AggregationOperator::as_str) + .to_vec(); + + // Spatial-of-temporal excludes topk (no topk(quantile_over_time(...)) pattern) + let spatial_ops_no_topk: Vec<&str> = [ + AggregationOperator::Sum, + AggregationOperator::Count, + AggregationOperator::Avg, + AggregationOperator::Quantile, + AggregationOperator::Min, + AggregationOperator::Max, + ] + .map(AggregationOperator::as_str) + .to_vec(); + // ONLY_TEMPORAL pattern 1: quantile_over_time(phi, metric[range]) let ot_quantile = PromQLPattern::new(PromQLPatternBuilder::function( - vec!["quantile_over_time"], + vec![PromQLFunction::QuantileOverTime.as_str()], vec![ PromQLPatternBuilder::number(None, None), range_vector_pattern(), @@ -20,15 +60,7 @@ pub fn build_patterns() -> Vec<(QueryPatternType, PromQLPattern)> { // ONLY_TEMPORAL pattern 2: sum_over_time/count_over_time/... (metric[range]) let ot_temporal_funcs = PromQLPattern::new(PromQLPatternBuilder::function( - vec![ - "sum_over_time", - "count_over_time", - "avg_over_time", - "min_over_time", - "max_over_time", - "increase", - "rate", - ], + temporal_funcs.clone(), vec![range_vector_pattern()], Some("function"), Some("function_args"), @@ -36,7 +68,7 @@ pub fn build_patterns() -> Vec<(QueryPatternType, PromQLPattern)> { // ONLY_SPATIAL pattern: agg_op(metric) let os_spatial = PromQLPattern::new(PromQLPatternBuilder::aggregation( - vec!["sum", "count", "avg", "quantile", "min", "max", "topk"], + spatial_ops, metric_pattern(), None, None, @@ -46,9 +78,9 @@ pub fn build_patterns() -> Vec<(QueryPatternType, PromQLPattern)> { // ONE_TEMPORAL_ONE_SPATIAL pattern 1: agg_op(quantile_over_time(phi, metric[range])) let ottos_quantile = PromQLPattern::new(PromQLPatternBuilder::aggregation( - vec!["sum", "count", "avg", "quantile", "min", "max"], + spatial_ops_no_topk.clone(), PromQLPatternBuilder::function( - vec!["quantile_over_time"], + vec![PromQLFunction::QuantileOverTime.as_str()], vec![ PromQLPatternBuilder::number(None, None), range_vector_pattern(), @@ -64,17 +96,9 @@ pub fn build_patterns() -> Vec<(QueryPatternType, PromQLPattern)> { // ONE_TEMPORAL_ONE_SPATIAL pattern 2: agg_op(temporal_func(metric[range])) let ottos_temporal = PromQLPattern::new(PromQLPatternBuilder::aggregation( - vec!["sum", "count", "avg", "quantile", "min", "max"], + spatial_ops_no_topk, PromQLPatternBuilder::function( - vec![ - "sum_over_time", - "count_over_time", - "avg_over_time", - "min_over_time", - "max_over_time", - "increase", - "rate", - ], + temporal_funcs, vec![range_vector_pattern()], Some("function"), Some("function_args"), diff --git a/asap-planner-rs/src/planner/single_query.rs b/asap-planner-rs/src/planner/single_query.rs index 221a7cb..f907aab 100644 --- a/asap-planner-rs/src/planner/single_query.rs +++ b/asap-planner-rs/src/planner/single_query.rs @@ -2,7 +2,9 @@ use asap_types::enums::{CleanupPolicy, WindowType}; use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; use promql_utilities::data_model::KeyByLabelNames; -use promql_utilities::query_logics::enums::{QueryPatternType, QueryTreatmentType, Statistic}; +use promql_utilities::query_logics::enums::{ + AggregationOperator, PromQLFunction, QueryPatternType, QueryTreatmentType, Statistic, +}; use promql_utilities::query_logics::logics::{ get_is_collapsable, map_statistic_to_precompute_operator, }; @@ -174,18 +176,15 @@ impl SingleQueryProcessor { match pattern_type { QueryPatternType::OnlyTemporal | QueryPatternType::OneTemporalOneSpatial => { let fn_name = match_result.get_function_name().unwrap_or_default(); - match fn_name.as_str() { - "quantile_over_time" | "sum_over_time" | "count_over_time" - | "avg_over_time" => QueryTreatmentType::Approximate, + match fn_name.parse::() { + Ok(f) if f.is_approximate() => QueryTreatmentType::Approximate, _ => QueryTreatmentType::Exact, } } QueryPatternType::OnlySpatial => { let op = match_result.get_aggregation_op().unwrap_or_default(); - match op.as_str() { - "quantile" | "sum" | "count" | "avg" | "topk" => { - QueryTreatmentType::Approximate - } + match op.parse::() { + Ok(o) if o.is_approximate() => QueryTreatmentType::Approximate, _ => QueryTreatmentType::Exact, } } @@ -246,12 +245,18 @@ impl SingleQueryProcessor { if pattern_type == QueryPatternType::OnlyTemporal { let fn_name = match_result.get_function_name().unwrap_or_default(); - if matches!(fn_name.as_str(), "rate" | "increase" | "quantile_over_time") { + let parsed_fn = fn_name.parse::(); + if matches!( + parsed_fn, + Ok(PromQLFunction::Rate + | PromQLFunction::Increase + | PromQLFunction::QuantileOverTime) + ) { let num_data_points = self.t_repeat as f64 / self.prometheus_scrape_interval as f64; if num_data_points < 60.0 { return false; } - if fn_name == "quantile_over_time" { + if parsed_fn == Ok(PromQLFunction::QuantileOverTime) { if let Some(range_dur) = match_result.get_range_duration() { let range_secs = range_dur.num_seconds() as f64; if range_secs / self.t_repeat as f64 > 15.0 { @@ -373,7 +378,12 @@ fn get_label_routing( QueryPatternType::OneTemporalOneSpatial => { let fn_name = match_result.get_function_name().unwrap_or_default(); let agg_op = match_result.get_aggregation_op().unwrap_or_default(); - if !get_is_collapsable(&fn_name, &agg_op) { + let collapsable = fn_name + .parse::() + .ok() + .zip(agg_op.parse::().ok()) + .is_some_and(|(f, o)| get_is_collapsable(f, o)); + if !collapsable { (KeyByLabelNames::empty(), all_labels.clone()) } else { let spatial_output = diff --git a/asap-query-engine/src/engines/simple_engine.rs b/asap-query-engine/src/engines/simple_engine.rs index 58f693c..83e23d8 100644 --- a/asap-query-engine/src/engines/simple_engine.rs +++ b/asap-query-engine/src/engines/simple_engine.rs @@ -9,6 +9,7 @@ use crate::engines::query_result::{InstantVectorElement, QueryResult, RangeVecto use crate::stores::{Store, TimestampedBucketsMap}; use core::panic; use promql_utilities::get_is_collapsable; +use promql_utilities::query_logics::enums::{AggregationOperator, PromQLFunction}; use serde_json::Value; use std::collections::HashMap; use std::sync::Arc; @@ -158,7 +159,7 @@ impl SimpleEngine { temporal_pattern_blocks.insert( "quantile".to_string(), PromQLPatternBuilder::function( - vec!["quantile_over_time"], + vec![PromQLFunction::QuantileOverTime.as_str()], vec![ PromQLPatternBuilder::number(None, Some("quantile_param")), PromQLPatternBuilder::matrix_selector( @@ -203,10 +204,31 @@ impl SimpleEngine { // Create spatial pattern blocks let mut spatial_pattern_blocks = HashMap::new(); + let spatial_ops_all: Vec<&str> = [ + AggregationOperator::Sum, + AggregationOperator::Count, + AggregationOperator::Avg, + AggregationOperator::Quantile, + AggregationOperator::Min, + AggregationOperator::Max, + AggregationOperator::Topk, + ] + .map(AggregationOperator::as_str) + .to_vec(); + let spatial_ops_no_topk: Vec<&str> = [ + AggregationOperator::Sum, + AggregationOperator::Count, + AggregationOperator::Avg, + AggregationOperator::Quantile, + AggregationOperator::Min, + AggregationOperator::Max, + ] + .map(AggregationOperator::as_str) + .to_vec(); spatial_pattern_blocks.insert( "generic".to_string(), PromQLPatternBuilder::aggregation( - vec!["sum", "count", "avg", "quantile", "min", "max", "topk"], + spatial_ops_all, PromQLPatternBuilder::metric(None, None, None, Some("metric")), None, None, @@ -230,19 +252,18 @@ impl SimpleEngine { PromQLPattern::new(blocks[pattern_type].clone()) } - fn spatial_of_temporal_pattern( - temporal_block: &Option>, - ) -> PromQLPattern { - let pattern = PromQLPatternBuilder::aggregation( - vec!["sum", "count", "avg", "quantile", "min", "max"], - temporal_block.clone(), - None, - None, - None, - Some("aggregation"), - ); - PromQLPattern::new(pattern) - } + let spatial_of_temporal_pattern = + |temporal_block: &Option>| -> PromQLPattern { + let pattern = PromQLPatternBuilder::aggregation( + spatial_ops_no_topk.clone(), + temporal_block.clone(), + None, + None, + None, + Some("aggregation"), + ); + PromQLPattern::new(pattern) + }; // Create controller patterns let mut controller_patterns = HashMap::new(); @@ -1089,9 +1110,15 @@ impl SimpleEngine { QueryPatternType::OneTemporalOneSpatial => { let temporal_aggregation = match_result.get_function_name().unwrap(); let spatial_aggregation = match_result.get_aggregation_op().unwrap(); - match get_is_collapsable(&temporal_aggregation, &spatial_aggregation) { - false => all_labels.clone(), - true => get_spatial_aggregation_output_labels(&match_result, &all_labels), + let collapsable = temporal_aggregation + .parse::() + .ok() + .zip(spatial_aggregation.parse::().ok()) + .is_some_and(|(f, o)| get_is_collapsable(f, o)); + if collapsable { + get_spatial_aggregation_output_labels(&match_result, &all_labels) + } else { + all_labels.clone() } } }; @@ -1585,13 +1612,10 @@ impl SimpleEngine { _ => query_data.aggregation_info.get_name().to_lowercase(), }; - let statistics: Vec = if statistic_name == "avg" { - vec![Statistic::Sum, Statistic::Count] - } else if let Ok(stat) = statistic_name.parse::() { - vec![stat] - } else { - vec![] - }; + let statistics: Vec = statistic_name + .parse::() + .map(|o| o.to_statistics()) + .unwrap_or_default(); let data_range_ms = match query_pattern_type { QueryPatternType::OnlySpatial => None, @@ -1850,13 +1874,10 @@ impl SimpleEngine { } }; - let statistics_to_compute: Vec = if statistic_name == "avg" { - vec![Statistic::Sum, Statistic::Count] - } else if let Ok(stat) = statistic_name.parse::() { - vec![stat] - } else { - panic!("Unsupported statistic: {}", statistic_name); - }; + let statistics_to_compute: Vec = statistic_name + .parse::() + .map(|o| o.to_statistics()) + .unwrap_or_else(|_| panic!("Unsupported statistic: {}", statistic_name)); if statistics_to_compute.len() != 1 { warn!( @@ -1975,13 +1996,10 @@ impl SimpleEngine { .get_name() .to_lowercase(); - let statistics_to_compute: Vec = if statistic_name == "avg" { - vec![Statistic::Sum, Statistic::Count] - } else if let Ok(stat) = statistic_name.parse::() { - vec![stat] - } else { - panic!("Unsupported statistic: {}", statistic_name); - }; + let statistics_to_compute: Vec = statistic_name + .parse::() + .map(|o| o.to_statistics()) + .unwrap_or_else(|_| panic!("Unsupported statistic: {}", statistic_name)); if statistics_to_compute.len() != 1 { warn!( @@ -2732,9 +2750,15 @@ impl SimpleEngine { let spatial_aggregation = match_result.get_aggregation_op().unwrap(); // iff temporal outer labels issubset of spatial inner labels, collapse // SQL issue: take into account labels from the query, not needed at present because only uses promql translations - match get_is_collapsable(&temporal_aggregation, &spatial_aggregation) { - false => all_labels.clone(), - true => get_spatial_aggregation_output_labels(&match_result, &all_labels), + let collapsable = temporal_aggregation + .parse::() + .ok() + .zip(spatial_aggregation.parse::().ok()) + .is_some_and(|(f, o)| get_is_collapsable(f, o)); + if collapsable { + get_spatial_aggregation_output_labels(&match_result, &all_labels) + } else { + all_labels.clone() } } }; From aa217ab222e160c3433c3a49a528c66565f08d28 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Thu, 9 Apr 2026 22:22:59 -0400 Subject: [PATCH 3/3] Changed accumulator strings to enums --- .../rs/asap_types/src/aggregation_config.rs | 19 +- .../rs/asap_types/src/capability_matching.rs | 58 +++--- .../dependencies/rs/asap_types/src/enums.rs | 3 + .../src/query_logics/enums.rs | 179 ++++++++++++++++-- .../src/query_logics/logics.rs | 59 +++--- .../src/query_logics/parsing.rs | 6 +- asap-planner-rs/src/output/generator.rs | 2 +- asap-planner-rs/src/planner/logics.rs | 30 +-- asap-planner-rs/src/planner/single_query.rs | 26 +-- .../src/planner/sql_single_query.rs | 4 +- .../benches/simple_store_bench.rs | 13 +- .../src/bin/bench_precompute_sketch.rs | 4 +- .../src/bin/show_logical_plans.rs | 56 +++--- .../src/bin/test_e2e_precompute.rs | 4 +- .../src/bin/test_offline_precomputes.rs | 22 +-- asap-query-engine/src/data_model/enums.rs | 1 + .../src/data_model/precomputed_output.rs | 30 +-- asap-query-engine/src/data_model/traits.rs | 4 +- .../src/engines/logical/plan_builder.rs | 128 +++++++------ .../src/engines/physical/accumulator_serde.rs | 3 +- .../physical/summary_merge_multiple_exec.rs | 4 +- .../src/engines/simple_engine.rs | 92 ++++----- .../src/engines/window_merger.rs | 10 +- .../precompute_engine/accumulator_factory.rs | 124 ++++++------ .../src/precompute_engine/worker.rs | 70 +++++-- .../count_min_sketch_accumulator.rs | 12 +- .../count_min_sketch_with_heap_accumulator.rs | 8 +- .../datasketches_kll_accumulator.rs | 11 +- .../delta_set_aggregator_accumulator.rs | 8 +- .../hydra_kll_accumulator.rs | 7 +- .../increase_accumulator.rs | 6 +- .../min_max_accumulator.rs | 8 +- .../multiple_increase_accumulator.rs | 8 +- .../multiple_min_max_accumulator.rs | 8 +- .../multiple_sum_accumulator.rs | 8 +- .../set_aggregator_accumulator.rs | 8 +- .../precompute_operators/sum_accumulator.rs | 8 +- .../src/stores/simple_map_store/global.rs | 7 +- .../stores/simple_map_store/legacy/global.rs | 5 +- .../stores/simple_map_store/legacy/per_key.rs | 5 +- .../src/stores/simple_map_store/per_key.rs | 10 +- .../src/tests/capability_matching_tests.rs | 56 +++++- .../datafusion/accumulator_serde_tests.rs | 17 +- .../datafusion/dispatch_arithmetic_tests.rs | 11 +- .../datafusion/plan_builder_binary_tests.rs | 6 +- .../plan_builder_regression_tests.rs | 47 ++--- .../plan_execution_arithmetic_tests.rs | 35 ++-- .../plan_execution_dual_input_tests.rs | 22 +-- .../plan_execution_temporal_tests.rs | 32 ++-- .../tests/datafusion/plan_execution_tests.rs | 40 ++-- .../datafusion/structural_matching_tests.rs | 11 +- .../src/tests/elastic_dsl_query_tests.rs | 12 +- .../src/tests/sql_pattern_matching_tests.rs | 6 +- .../src/tests/store_correctness_tests.rs | 38 ++-- .../src/tests/test_utilities/comparison.rs | 6 +- .../tests/test_utilities/config_builders.rs | 10 +- .../tests/test_utilities/engine_factories.rs | 44 ++--- .../tests/e2e_precompute_equivalence.rs | 10 +- 58 files changed, 893 insertions(+), 588 deletions(-) diff --git a/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs b/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs index 6820d4b..aae18a7 100644 --- a/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs @@ -7,11 +7,12 @@ use crate::enums::{QueryLanguage, WindowType}; use crate::traits::SerializableToSink; use crate::utils::normalize_spatial_filter; use promql_utilities::data_model::KeyByLabelNames; +use promql_utilities::query_logics::enums::AggregationType; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AggregationConfig { pub aggregation_id: u64, - pub aggregation_type: String, + pub aggregation_type: AggregationType, pub aggregation_sub_type: String, pub parameters: HashMap, pub grouping_labels: KeyByLabelNames, @@ -41,8 +42,8 @@ pub struct AggregationConfig { pub struct AggregationIdInfo { pub aggregation_id_for_key: u64, pub aggregation_id_for_value: u64, - pub aggregation_type_for_key: String, - pub aggregation_type_for_value: String, + pub aggregation_type_for_key: AggregationType, + pub aggregation_type_for_value: AggregationType, } // TODO: need to implement deserialization methods @@ -51,7 +52,7 @@ impl AggregationConfig { #[allow(clippy::too_many_arguments)] pub fn new( aggregation_id: u64, - aggregation_type: String, + aggregation_type: AggregationType, aggregation_sub_type: String, parameters: HashMap, grouping_labels: KeyByLabelNames, @@ -116,10 +117,11 @@ impl AggregationConfig { .as_u64() .ok_or("Missing aggregationId")?; - let aggregation_type = data["aggregationType"] + let aggregation_type: AggregationType = data["aggregationType"] .as_str() .ok_or("Missing aggregationType")? - .to_string(); + .parse() + .map_err(|e: String| e)?; let aggregation_sub_type = data["aggregationSubType"] .as_str() @@ -241,10 +243,11 @@ impl AggregationConfig { .collect(), ); - let aggregation_type = aggregation_data["aggregationType"] + let aggregation_type: AggregationType = aggregation_data["aggregationType"] .as_str() .ok_or_else(|| anyhow::anyhow!("Missing aggregationType"))? - .to_string(); + .parse() + .map_err(|e: String| anyhow::anyhow!(e))?; let aggregation_sub_type = aggregation_data["aggregationSubType"] .as_str() diff --git a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs index 2dbb1a2..570c085 100644 --- a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs +++ b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs @@ -9,29 +9,32 @@ use crate::aggregation_config::{AggregationConfig, AggregationIdInfo}; use crate::enums::WindowType; use crate::query_requirements::QueryRequirements; use crate::utils::normalize_spatial_filter; +use promql_utilities::query_logics::enums::AggregationType; // --------------------------------------------------------------------------- // Pure compatibility helpers // --------------------------------------------------------------------------- -/// Returns the aggregation_type strings that can serve this statistic. -pub fn compatible_agg_types(stat: Statistic) -> &'static [&'static str] { +/// Returns the aggregation types that can serve this statistic. +pub fn compatible_agg_types(stat: Statistic) -> &'static [AggregationType] { match stat { - Statistic::Sum => &["Sum", "MultipleSumAccumulator"], + Statistic::Sum => &[AggregationType::Sum, AggregationType::MultipleSum], Statistic::Count => &[ - "CountMinSketch", - "CountMinSketchWithHeap", - "CountMinSketchWithHeapAccumulator", + AggregationType::CountMinSketch, + AggregationType::CountMinSketchWithHeap, ], - Statistic::Min => &["MinMax", "MultipleMinMaxAccumulator"], - Statistic::Max => &["MinMax", "MultipleMinMaxAccumulator"], - Statistic::Quantile => &["DatasketchesKLL", "HydraKLL"], - Statistic::Rate | Statistic::Increase => &["Increase", "MultipleIncreaseAccumulator"], - Statistic::Cardinality => &["SetAggregator", "DeltaSetAggregator"], - Statistic::Topk => &[ - "CountMinSketchWithHeap", - "CountMinSketchWithHeapAccumulator", + Statistic::Min | Statistic::Max => { + &[AggregationType::MinMax, AggregationType::MultipleMinMax] + } + Statistic::Quantile => &[AggregationType::DatasketchesKLL, AggregationType::HydraKLL], + Statistic::Rate | Statistic::Increase => { + &[AggregationType::Increase, AggregationType::MultipleIncrease] + } + Statistic::Cardinality => &[ + AggregationType::SetAggregator, + AggregationType::DeltaSetAggregator, ], + Statistic::Topk => &[AggregationType::CountMinSketchWithHeap], } } @@ -47,20 +50,13 @@ pub fn required_sub_type(stat: Statistic) -> Option<&'static str> { /// Whether this value aggregation type requires a paired key aggregation /// (`SetAggregator` or `DeltaSetAggregator`). -pub fn is_multi_population_value_type(agg_type: &str) -> bool { - matches!( - agg_type, - "MultipleSumAccumulator" - | "MultipleMinMaxAccumulator" - | "MultipleIncreaseAccumulator" - | "CountMinSketchWithHeap" - | "CountMinSketchWithHeapAccumulator" - ) +pub fn is_multi_population_value_type(agg_type: AggregationType) -> bool { + agg_type.is_multi_population_value_type() } /// Whether this type is a key aggregation (tracks which label-value combinations exist). -fn is_key_agg_type(agg_type: &str) -> bool { - matches!(agg_type, "SetAggregator" | "DeltaSetAggregator") +fn is_key_agg_type(agg_type: AggregationType) -> bool { + agg_type.is_key_agg_type() } /// Window compatibility: can `config` serve a query needing `data_range_ms`? @@ -153,7 +149,7 @@ pub fn find_compatible_aggregation( .values() .filter(|c| { let ok = c.metric == requirements.metric - && types.contains(&c.aggregation_type.as_str()) + && types.contains(&c.aggregation_type) && sub_type.is_none_or(|st| c.aggregation_sub_type == st) && window_compatible(c, requirements.data_range_ms) && labels_compatible(&c.grouping_labels, &requirements.grouping_labels) @@ -219,11 +215,11 @@ pub fn find_compatible_aggregation( } // If value type is multi-population, find the paired key aggregation. - let key_agg: &AggregationConfig = if is_multi_population_value_type(&value_agg.aggregation_type) + let key_agg: &AggregationConfig = if is_multi_population_value_type(value_agg.aggregation_type) { let ka = configs .values() - .find(|c| c.metric == requirements.metric && is_key_agg_type(&c.aggregation_type)); + .find(|c| c.metric == requirements.metric && is_key_agg_type(c.aggregation_type)); if ka.is_none() { warn!( metric = %requirements.metric, @@ -247,9 +243,9 @@ pub fn find_compatible_aggregation( Some(AggregationIdInfo { aggregation_id_for_value: value_agg.aggregation_id, - aggregation_type_for_value: value_agg.aggregation_type.clone(), + aggregation_type_for_value: value_agg.aggregation_type, aggregation_id_for_key: key_agg.aggregation_id, - aggregation_type_for_key: key_agg.aggregation_type.clone(), + aggregation_type_for_key: key_agg.aggregation_type, }) } @@ -280,7 +276,7 @@ mod tests { let spatial_filter_normalized = normalize_spatial_filter(spatial_filter); AggregationConfig { aggregation_id: id, - aggregation_type: agg_type.to_string(), + aggregation_type: agg_type.parse::().expect("valid agg type"), aggregation_sub_type: sub_type.to_string(), parameters: HashMap::new(), grouping_labels, diff --git a/asap-common/dependencies/rs/asap_types/src/enums.rs b/asap-common/dependencies/rs/asap_types/src/enums.rs index c20e88b..7d511c3 100644 --- a/asap-common/dependencies/rs/asap_types/src/enums.rs +++ b/asap-common/dependencies/rs/asap_types/src/enums.rs @@ -1,6 +1,9 @@ use std::fmt; use std::str::FromStr; +// Re-export AggregationType from promql_utilities (defined there to avoid circular deps). +pub use promql_utilities::query_logics::enums::AggregationType; + #[derive(clap::ValueEnum, Clone, Copy, Debug, PartialEq)] #[allow(non_camel_case_types)] pub enum QueryLanguage { diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs index 883c3a6..05b3435 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs @@ -1,4 +1,6 @@ use serde::{Deserialize, Serialize}; +use std::fmt; +use std::str::FromStr; use tracing::debug; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -212,6 +214,34 @@ impl AggregationOperator { } } + pub fn as_str_slice() -> &'static [&'static str] { + &["sum", "count", "avg", "quantile", "min", "max", "topk"] + } +} + +impl fmt::Display for AggregationOperator { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +impl FromStr for AggregationOperator { + type Err = String; + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "sum" => Ok(AggregationOperator::Sum), + "count" => Ok(AggregationOperator::Count), + "avg" => Ok(AggregationOperator::Avg), + "quantile" => Ok(AggregationOperator::Quantile), + "min" => Ok(AggregationOperator::Min), + "max" => Ok(AggregationOperator::Max), + "topk" => Ok(AggregationOperator::Topk), + other => Err(format!("Unknown aggregation operator: '{other}'")), + } + } +} + +impl AggregationOperator { /// Returns `true` for operators whose result requires approximate pre-aggregation. pub fn is_approximate(self) -> bool { matches!( @@ -225,28 +255,155 @@ impl AggregationOperator { } } -impl std::fmt::Display for AggregationOperator { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { +/// Concrete aggregation/sketch type used in precompute configs and accumulator dispatch. +/// +/// `Display` outputs the canonical PascalCase name used in YAML/JSON configs. +/// `FromStr` accepts the canonical name plus legacy aliases (e.g. "KLL" → `DatasketchesKLL`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum AggregationType { + // ---------- single-population (non-keyed) ---------- + Sum, + Increase, + MinMax, + DatasketchesKLL, + // ---------- multi-population (keyed) ---------- + MultipleSum, + MultipleIncrease, + MultipleMinMax, + HydraKLL, + CountMinSketch, + CountMinSketchWithHeap, + // ---------- cardinality / set tracking ---------- + SetAggregator, + DeltaSetAggregator, + HLL, + // ---------- legacy config wrapper names ---------- + SingleSubpopulation, + MultipleSubpopulation, +} + +impl AggregationType { + pub fn as_str(self) -> &'static str { + match self { + AggregationType::Sum => "Sum", + AggregationType::Increase => "Increase", + AggregationType::MinMax => "MinMax", + AggregationType::DatasketchesKLL => "DatasketchesKLL", + AggregationType::MultipleSum => "MultipleSum", + AggregationType::MultipleIncrease => "MultipleIncrease", + AggregationType::MultipleMinMax => "MultipleMinMax", + AggregationType::HydraKLL => "HydraKLL", + AggregationType::CountMinSketch => "CountMinSketch", + AggregationType::CountMinSketchWithHeap => "CountMinSketchWithHeap", + AggregationType::SetAggregator => "SetAggregator", + AggregationType::DeltaSetAggregator => "DeltaSetAggregator", + AggregationType::HLL => "HLL", + AggregationType::SingleSubpopulation => "SingleSubpopulation", + AggregationType::MultipleSubpopulation => "MultipleSubpopulation", + } + } + + /// Returns `true` if this type produces keyed (multi-population) accumulators. + pub fn is_keyed(self) -> bool { + matches!( + self, + AggregationType::MultipleSubpopulation + | AggregationType::MultipleSum + | AggregationType::MultipleIncrease + | AggregationType::MultipleMinMax + | AggregationType::CountMinSketch + | AggregationType::CountMinSketchWithHeap + | AggregationType::HydraKLL + ) + } + + /// Returns `true` if this type needs a paired key aggregation (SetAggregator / DeltaSetAggregator). + pub fn is_multi_population_value_type(self) -> bool { + matches!( + self, + AggregationType::MultipleSum + | AggregationType::MultipleMinMax + | AggregationType::MultipleIncrease + | AggregationType::CountMinSketch + | AggregationType::CountMinSketchWithHeap + ) + } + + /// Returns `true` if this is a key-tracking aggregation type. + pub fn is_key_agg_type(self) -> bool { + matches!( + self, + AggregationType::SetAggregator | AggregationType::DeltaSetAggregator + ) + } +} + +impl fmt::Display for AggregationType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str(self.as_str()) } } -impl std::str::FromStr for AggregationOperator { +impl FromStr for AggregationType { type Err = String; + fn from_str(s: &str) -> Result { match s { - "sum" => Ok(AggregationOperator::Sum), - "count" => Ok(AggregationOperator::Count), - "avg" => Ok(AggregationOperator::Avg), - "quantile" => Ok(AggregationOperator::Quantile), - "min" => Ok(AggregationOperator::Min), - "max" => Ok(AggregationOperator::Max), - "topk" => Ok(AggregationOperator::Topk), - other => Err(format!("Unknown aggregation operator: '{other}'")), + // Canonical names + "Sum" => Ok(AggregationType::Sum), + "Increase" => Ok(AggregationType::Increase), + "MinMax" => Ok(AggregationType::MinMax), + "DatasketchesKLL" => Ok(AggregationType::DatasketchesKLL), + "MultipleSum" => Ok(AggregationType::MultipleSum), + "MultipleIncrease" => Ok(AggregationType::MultipleIncrease), + "MultipleMinMax" => Ok(AggregationType::MultipleMinMax), + "HydraKLL" => Ok(AggregationType::HydraKLL), + "CountMinSketch" => Ok(AggregationType::CountMinSketch), + "CountMinSketchWithHeap" => Ok(AggregationType::CountMinSketchWithHeap), + "SetAggregator" => Ok(AggregationType::SetAggregator), + "DeltaSetAggregator" => Ok(AggregationType::DeltaSetAggregator), + "HLL" | "HyperLogLog" => Ok(AggregationType::HLL), + "SingleSubpopulation" => Ok(AggregationType::SingleSubpopulation), + "MultipleSubpopulation" => Ok(AggregationType::MultipleSubpopulation), + // Legacy accumulator-suffixed aliases + "SumAccumulator" | "SumAggregator" | "sum" => Ok(AggregationType::Sum), + "IncreaseAccumulator" | "IncreaseAggregator" | "increase" => { + Ok(AggregationType::Increase) + } + "MinMaxAccumulator" | "MinMaxAggregator" | "min_max" => Ok(AggregationType::MinMax), + "DatasketchesKLLAccumulator" | "KLL" | "kll" | "datasketches_kll" => { + Ok(AggregationType::DatasketchesKLL) + } + "MultipleSumAccumulator" | "multiple_sum" => Ok(AggregationType::MultipleSum), + "MultipleIncreaseAccumulator" | "multiple_increase" => { + Ok(AggregationType::MultipleIncrease) + } + "MultipleMinMaxAccumulator" | "multiple_min_max" => Ok(AggregationType::MultipleMinMax), + "HydraKllSketchAccumulator" | "hydra_kll" => Ok(AggregationType::HydraKLL), + "CountMinSketchAccumulator" | "CMS" | "cms" | "count_min_sketch" => { + Ok(AggregationType::CountMinSketch) + } + "CountMinSketchWithHeapAccumulator" => Ok(AggregationType::CountMinSketchWithHeap), + "SetAggregatorAccumulator" => Ok(AggregationType::SetAggregator), + "DeltaSetAggregatorAccumulator" => Ok(AggregationType::DeltaSetAggregator), + _ => Err(format!("Unknown aggregation type: '{s}'")), } } } +impl Serialize for AggregationType { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_str(self.as_str()) + } +} + +impl<'de> Deserialize<'de> for AggregationType { + fn deserialize>(deserializer: D) -> Result { + let s = String::deserialize(deserializer)?; + s.parse().map_err(serde::de::Error::custom) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs index 0684605..d9eb4ca 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs @@ -1,5 +1,5 @@ use crate::query_logics::enums::{ - AggregationOperator, PromQLFunction, QueryTreatmentType, Statistic, + AggregationOperator, AggregationType, PromQLFunction, QueryTreatmentType, Statistic, }; use tracing::debug; @@ -8,7 +8,7 @@ use tracing::debug; pub fn map_statistic_to_precompute_operator( statistic: Statistic, treatment_type: QueryTreatmentType, -) -> Result<(String, String), String> { +) -> Result<(AggregationType, String), String> { debug!( "Mapping statistic {:?} with treatment type {:?} to precompute operator", statistic, treatment_type @@ -18,17 +18,17 @@ pub fn map_statistic_to_precompute_operator( if treatment_type == QueryTreatmentType::Exact { Err("Statistic Quantile cannot be computed exactly".to_string()) } else { - Ok(("DatasketchesKLL".to_string(), "".to_string())) - //Ok(("HydraKLL".to_string(), "".to_string())) + Ok((AggregationType::DatasketchesKLL, "".to_string())) + //Ok((AggregationType::HydraKLL, "".to_string())) } } Statistic::Min | Statistic::Max => { if treatment_type == QueryTreatmentType::Approximate { - Ok(("DatasketchesKLL".to_string(), "".to_string())) - //Ok(("HydraKLL".to_string(), "".to_string())) + Ok((AggregationType::DatasketchesKLL, "".to_string())) + //Ok((AggregationType::HydraKLL, "".to_string())) } else { Ok(( - "MultipleMinMax".to_string(), + AggregationType::MultipleMinMax, statistic.to_string().to_lowercase(), )) } @@ -36,20 +36,20 @@ pub fn map_statistic_to_precompute_operator( Statistic::Sum | Statistic::Count => { if treatment_type == QueryTreatmentType::Approximate { Ok(( - "CountMinSketch".to_string(), + AggregationType::CountMinSketch, statistic.to_string().to_lowercase(), )) } else { Ok(( - "MultipleSum".to_string(), + AggregationType::MultipleSum, statistic.to_string().to_lowercase(), )) } } Statistic::Rate | Statistic::Increase => { - Ok(("MultipleIncrease".to_string(), "".to_string())) + Ok((AggregationType::MultipleIncrease, "".to_string())) } - Statistic::Topk => Ok(("CountMinSketchWithHeap".to_string(), "topk".to_string())), + Statistic::Topk => Ok((AggregationType::CountMinSketchWithHeap, "topk".to_string())), _ => Err(format!("Statistic {statistic:?} not supported")), } } @@ -57,7 +57,7 @@ pub fn map_statistic_to_precompute_operator( /// Check if a precompute operator supports subpopulations (multiple keys) pub fn does_precompute_operator_support_subpopulations( statistic: Statistic, - precompute_operator: &str, + precompute_operator: AggregationType, ) -> bool { debug!( "Checking if precompute operator '{}' supports subpopulations for statistic {:?}", @@ -65,17 +65,22 @@ pub fn does_precompute_operator_support_subpopulations( ); match precompute_operator { // Single-key operators - "Increase" | "MinMax" | "Sum" | "DatasketchesKLL" => false, + AggregationType::Increase + | AggregationType::MinMax + | AggregationType::Sum + | AggregationType::DatasketchesKLL => false, // Multi-key operators - "MultipleIncrease" | "MultipleMinMax" | "MultipleSum" | "HydraKLL" => true, + AggregationType::MultipleIncrease + | AggregationType::MultipleMinMax + | AggregationType::MultipleSum + | AggregationType::HydraKLL => true, // CountMinSketch supports subpopulations only for certain statistics - "CountMinSketch" => matches!(statistic, Statistic::Sum | Statistic::Count), + AggregationType::CountMinSketch => matches!(statistic, Statistic::Sum | Statistic::Count), - // "CountMinSketchWithHeap" is only supported for Topk - // Other usages of CountMinSketchWithHeap will fall through. - "CountMinSketchWithHeap" if matches!(statistic, Statistic::Topk) => false, + // CountMinSketchWithHeap is only supported for Topk — does not support subpopulations + AggregationType::CountMinSketchWithHeap if matches!(statistic, Statistic::Topk) => false, // Default: not supported _ => panic!("Unexpected precompute operator: {}", precompute_operator), @@ -114,13 +119,13 @@ mod tests { let result = map_statistic_to_precompute_operator(Statistic::Sum, QueryTreatmentType::Exact) .unwrap(); - assert_eq!(result, ("MultipleSum".to_string(), "sum".to_string())); + assert_eq!(result, (AggregationType::MultipleSum, "sum".to_string())); // Test approximate sum let result = map_statistic_to_precompute_operator(Statistic::Sum, QueryTreatmentType::Approximate) .unwrap(); - assert_eq!(result, ("CountMinSketch".to_string(), "sum".to_string())); + assert_eq!(result, (AggregationType::CountMinSketch, "sum".to_string())); // Test exact quantile (should fail) let result = @@ -133,8 +138,8 @@ mod tests { QueryTreatmentType::Approximate, ) .unwrap(); - assert_eq!(result, ("DatasketchesKLL".to_string(), "".to_string())); - //assert_eq!(result, ("HydraKLL".to_string(), "".to_string())); + assert_eq!(result, (AggregationType::DatasketchesKLL, "".to_string())); + //assert_eq!(result, (AggregationType::HydraKLL, "".to_string())); } #[test] @@ -142,25 +147,25 @@ mod tests { // Test MultipleSum supports subpopulations assert!(does_precompute_operator_support_subpopulations( Statistic::Sum, - "MultipleSum" + AggregationType::MultipleSum, )); // Test DatasketchesKLL does not support subpopulations assert!(!does_precompute_operator_support_subpopulations( Statistic::Quantile, - "DatasketchesKLL" + AggregationType::DatasketchesKLL, )); // Test HydraKLL supports subpopulations assert!(does_precompute_operator_support_subpopulations( Statistic::Quantile, - "HydraKLL" + AggregationType::HydraKLL, )); // Test CountMinSketch with valid statistic assert!(does_precompute_operator_support_subpopulations( Statistic::Sum, - "CountMinSketch" + AggregationType::CountMinSketch, )); } @@ -171,7 +176,7 @@ mod tests { .unwrap(); assert_eq!( result, - ("CountMinSketchWithHeap".to_string(), "topk".to_string()) + (AggregationType::CountMinSketchWithHeap, "topk".to_string()) ); } diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs index b57c699..e5bb875 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs @@ -6,7 +6,7 @@ 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::{QueryPatternType, Statistic}; +use crate::query_logics::enums::{AggregationOperator, QueryPatternType, Statistic}; pub fn get_metric_and_spatial_filter(match_result: &PromQLMatchResult) -> (String, String) { debug!("Extracting metric and spatial filter from match result"); @@ -73,7 +73,7 @@ pub fn get_statistics_to_compute( if let Some(statistic_to_compute) = statistic_to_compute { debug!("Found statistic to compute: {}", statistic_to_compute); - if statistic_to_compute == "avg" { + if statistic_to_compute.parse::() == Ok(AggregationOperator::Avg) { vec![Statistic::Sum, Statistic::Count] } else if let Ok(stat) = statistic_to_compute.parse::() { vec![stat] @@ -101,7 +101,7 @@ pub fn get_spatial_aggregation_output_labels( .expect("aggregation token missing"); // Patching: When the query is topk, we should always return all labels - if aggregation_token.op.to_lowercase() == "topk" { + if aggregation_token.op.parse::() == Ok(AggregationOperator::Topk) { debug!("Aggregation operation is 'topk', returning all labels"); return all_labels.clone(); } diff --git a/asap-planner-rs/src/output/generator.rs b/asap-planner-rs/src/output/generator.rs index 69b168f..6da86ab 100644 --- a/asap-planner-rs/src/output/generator.rs +++ b/asap-planner-rs/src/output/generator.rs @@ -211,7 +211,7 @@ pub fn build_aggregation_entry(id: u32, cfg: &IntermediateAggConfig) -> YamlValu ); map.insert( YamlValue::String("aggregationType".to_string()), - YamlValue::String(cfg.aggregation_type.clone()), + YamlValue::String(cfg.aggregation_type.to_string()), ); let mut labels_map = serde_yaml::Mapping::new(); diff --git a/asap-planner-rs/src/planner/logics.rs b/asap-planner-rs/src/planner/logics.rs index 10dc0e9..3e415df 100644 --- a/asap-planner-rs/src/planner/logics.rs +++ b/asap-planner-rs/src/planner/logics.rs @@ -2,7 +2,7 @@ use crate::config::input::SketchParameterOverrides; use asap_types::enums::{CleanupPolicy, WindowType}; 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::{AggregationType, QueryPatternType, Statistic}; use promql_utilities::query_logics::logics::does_precompute_operator_support_subpopulations; use std::collections::HashMap; @@ -84,16 +84,22 @@ pub struct IntermediateWindowConfig { /// from the `topk(k, …)` query argument; SQL passes `None` (SQL never produces /// this operator today, so the `None` branch is unreachable in practice). pub fn build_sketch_parameters( - aggregation_type: &str, + aggregation_type: AggregationType, aggregation_sub_type: &str, topk_k: Option, sketch_params: Option<&SketchParameterOverrides>, ) -> Result, String> { match aggregation_type { - "Increase" | "MinMax" | "Sum" | "MultipleIncrease" | "MultipleMinMax" | "MultipleSum" - | "DeltaSetAggregator" | "SetAggregator" => Ok(HashMap::new()), - - "CountMinSketch" => { + AggregationType::Increase + | AggregationType::MinMax + | AggregationType::Sum + | AggregationType::MultipleIncrease + | AggregationType::MultipleMinMax + | AggregationType::MultipleSum + | AggregationType::DeltaSetAggregator + | AggregationType::SetAggregator => Ok(HashMap::new()), + + AggregationType::CountMinSketch => { let depth = sketch_params .and_then(|p| p.count_min_sketch.as_ref()) .map(|p| p.depth) @@ -108,7 +114,7 @@ pub fn build_sketch_parameters( Ok(m) } - "CountMinSketchWithHeap" => { + AggregationType::CountMinSketchWithHeap => { if aggregation_sub_type != "topk" { return Err(format!( "Aggregation sub-type {} for CountMinSketchWithHeap not supported", @@ -139,7 +145,7 @@ pub fn build_sketch_parameters( Ok(m) } - "DatasketchesKLL" => { + AggregationType::DatasketchesKLL => { let k = sketch_params .and_then(|p| p.datasketches_kll.as_ref()) .map(|p| p.k) @@ -149,7 +155,7 @@ pub fn build_sketch_parameters( Ok(m) } - "HydraKLL" => { + AggregationType::HydraKLL => { let row_num = sketch_params .and_then(|p| p.hydra_kll.as_ref()) .map(|p| p.row_num) @@ -182,12 +188,12 @@ pub fn build_sketch_parameters( /// PromQL wrapper: extracts the topk `k` from the match result when needed, /// then delegates to `build_sketch_parameters`. pub fn build_sketch_parameters_from_promql( - aggregation_type: &str, + aggregation_type: AggregationType, aggregation_sub_type: &str, match_result: &PromQLMatchResult, sketch_params: Option<&SketchParameterOverrides>, ) -> Result, String> { - let topk_k = if aggregation_type == "CountMinSketchWithHeap" { + let topk_k = if aggregation_type == AggregationType::CountMinSketchWithHeap { let k: u64 = match_result .tokens .get("aggregation") @@ -274,7 +280,7 @@ pub fn get_cleanup_param( pub fn set_subpopulation_labels( statistic: Statistic, - aggregation_type: &str, + aggregation_type: AggregationType, subpopulation_labels: &KeyByLabelNames, rollup_labels: &mut KeyByLabelNames, grouping_labels: &mut KeyByLabelNames, diff --git a/asap-planner-rs/src/planner/single_query.rs b/asap-planner-rs/src/planner/single_query.rs index f907aab..9bdb225 100644 --- a/asap-planner-rs/src/planner/single_query.rs +++ b/asap-planner-rs/src/planner/single_query.rs @@ -3,7 +3,8 @@ use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::{ - AggregationOperator, PromQLFunction, QueryPatternType, QueryTreatmentType, Statistic, + AggregationOperator, AggregationType, PromQLFunction, QueryPatternType, QueryTreatmentType, + Statistic, }; use promql_utilities::query_logics::logics::{ get_is_collapsable, map_statistic_to_precompute_operator, @@ -57,7 +58,7 @@ fn strip_parens(expr: &promql_parser::parser::Expr) -> &promql_parser::parser::E /// Internal representation of an aggregation config before IDs are assigned #[derive(Debug, Clone)] pub struct IntermediateAggConfig { - pub aggregation_type: String, + pub aggregation_type: AggregationType, pub aggregation_sub_type: String, pub window_type: WindowType, pub window_size: u64, @@ -315,7 +316,7 @@ impl SingleQueryProcessor { None, None, &spatial_filter, - |agg_type, agg_sub_type| { + |agg_type: AggregationType, agg_sub_type: &str| { build_sketch_parameters_from_promql( agg_type, agg_sub_type, @@ -410,7 +411,7 @@ pub fn build_agg_configs_for_statistics( table_name: Option<&str>, value_column: Option<&str>, spatial_filter: &str, - get_params: impl Fn(&str, &str) -> Result, String>, + get_params: impl Fn(AggregationType, &str) -> Result, String>, ) -> Result, String> { let mut configs = Vec::new(); @@ -422,17 +423,20 @@ pub fn build_agg_configs_for_statistics( let mut aggregated = KeyByLabelNames::empty(); set_subpopulation_labels( statistic, - &agg_type, + agg_type, subpopulation_labels, &mut rollup.clone(), &mut grouping, &mut aggregated, ); - if matches!(agg_type.as_str(), "CountMinSketch" | "HydraKLL") { - let delta_params = get_params("DeltaSetAggregator", "")?; + if matches!( + agg_type, + AggregationType::CountMinSketch | AggregationType::HydraKLL + ) { + let delta_params = get_params(AggregationType::DeltaSetAggregator, "")?; configs.push(IntermediateAggConfig { - aggregation_type: "DeltaSetAggregator".to_string(), + aggregation_type: AggregationType::DeltaSetAggregator, aggregation_sub_type: String::new(), window_type: window_cfg.window_type, window_size: window_cfg.window_size, @@ -448,7 +452,7 @@ pub fn build_agg_configs_for_statistics( }); } - let parameters = get_params(&agg_type, &agg_sub_type)?; + let parameters = get_params(agg_type, &agg_sub_type)?; configs.push(IntermediateAggConfig { aggregation_type: agg_type, aggregation_sub_type: agg_sub_type, @@ -477,7 +481,7 @@ mod tests { fn base_config() -> IntermediateAggConfig { IntermediateAggConfig { - aggregation_type: "MultipleIncrease".to_string(), + aggregation_type: AggregationType::MultipleIncrease, aggregation_sub_type: "rate".to_string(), window_type: WindowType::Tumbling, window_size: 300, @@ -511,7 +515,7 @@ mod tests { fn different_aggregation_type_produces_different_key() { let cfg1 = base_config(); let mut cfg2 = base_config(); - cfg2.aggregation_type = "DatasketchesKLL".to_string(); + cfg2.aggregation_type = AggregationType::DatasketchesKLL; assert_ne!(cfg1.identifying_key(), cfg2.identifying_key()); } diff --git a/asap-planner-rs/src/planner/sql_single_query.rs b/asap-planner-rs/src/planner/sql_single_query.rs index 53edbdd..057b320 100644 --- a/asap-planner-rs/src/planner/sql_single_query.rs +++ b/asap-planner-rs/src/planner/sql_single_query.rs @@ -2,7 +2,7 @@ use std::collections::HashSet; use asap_types::enums::{CleanupPolicy, WindowType}; use promql_utilities::data_model::KeyByLabelNames; -use promql_utilities::query_logics::enums::{QueryTreatmentType, Statistic}; +use promql_utilities::query_logics::enums::{AggregationType, QueryTreatmentType, Statistic}; use sql_utilities::ast_matching::sqlhelper::Table; use sql_utilities::ast_matching::sqlpattern_matcher::{QueryType, SQLPatternMatcher}; use sql_utilities::ast_matching::sqlpattern_parser::SQLPatternParser; @@ -120,7 +120,7 @@ impl SQLSingleQueryProcessor { Some(table_name), Some(&value_column), "", - |agg_type, agg_sub_type| { + |agg_type: AggregationType, agg_sub_type: &str| { build_sketch_parameters( agg_type, agg_sub_type, diff --git a/asap-query-engine/benches/simple_store_bench.rs b/asap-query-engine/benches/simple_store_bench.rs index 27a907c..d225188 100644 --- a/asap-query-engine/benches/simple_store_bench.rs +++ b/asap-query-engine/benches/simple_store_bench.rs @@ -27,7 +27,8 @@ use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; use promql_utilities::data_model::KeyByLabelNames; use query_engine_rust::data_model::{ - AggregateCore, CleanupPolicy, KeyByLabelValues, LockStrategy, StreamingConfig, WindowType, + AggregateCore, AggregationType, CleanupPolicy, KeyByLabelValues, LockStrategy, StreamingConfig, + WindowType, }; use query_engine_rust::precompute_operators::{DatasketchesKLLAccumulator, SumAccumulator}; use query_engine_rust::stores::simple_map_store::legacy::{ @@ -97,10 +98,10 @@ impl AccumulatorKind { } } - fn aggregation_type(self) -> &'static str { + fn aggregation_type(self) -> AggregationType { match self { - Self::Sum => "Sum", - Self::Kll => "DatasketchesKLL", + Self::Sum => AggregationType::Sum, + Self::Kll => AggregationType::DatasketchesKLL, } } @@ -120,14 +121,14 @@ impl AccumulatorKind { fn make_agg_config( agg_id: u64, - aggregation_type: &str, + aggregation_type: AggregationType, metric: &str, num_aggregates_to_retain: Option, read_count_threshold: Option, ) -> AggregationConfig { AggregationConfig::new( agg_id, - aggregation_type.to_string(), + aggregation_type, "".to_string(), HashMap::new(), KeyByLabelNames::empty(), diff --git a/asap-query-engine/src/bin/bench_precompute_sketch.rs b/asap-query-engine/src/bin/bench_precompute_sketch.rs index f99cfb4..16f2e66 100644 --- a/asap-query-engine/src/bin/bench_precompute_sketch.rs +++ b/asap-query-engine/src/bin/bench_precompute_sketch.rs @@ -1,5 +1,5 @@ use asap_types::aggregation_config::AggregationConfig; -use asap_types::enums::WindowType; +use asap_types::enums::{AggregationType, WindowType}; use clap::Parser; use prost::Message; use query_engine_rust::data_model::{ @@ -127,7 +127,7 @@ fn make_kll_streaming_config( let agg_config = AggregationConfig::new( aggregation_id, - "DatasketchesKLL".to_string(), + AggregationType::DatasketchesKLL, String::new(), params, promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), diff --git a/asap-query-engine/src/bin/show_logical_plans.rs b/asap-query-engine/src/bin/show_logical_plans.rs index 24b31ad..39bd02b 100644 --- a/asap-query-engine/src/bin/show_logical_plans.rs +++ b/asap-query-engine/src/bin/show_logical_plans.rs @@ -14,7 +14,7 @@ use datafusion::logical_expr::LogicalPlan; use datafusion_summary_library::{PrecomputedSummaryRead, SummaryInfer, SummaryMergeMultiple}; use promql_utilities::data_model::KeyByLabelNames; -use promql_utilities::query_logics::enums::Statistic; +use promql_utilities::query_logics::enums::{AggregationType, Statistic}; use query_engine_rust::data_model::AggregationIdInfo; use query_engine_rust::engines::simple_engine::{ QueryExecutionContext, QueryMetadata, StoreQueryParams, StoreQueryPlan, @@ -33,8 +33,8 @@ fn build_context( query_output_labels: Vec<&str>, grouping_labels: Vec<&str>, aggregated_labels: Vec<&str>, - agg_type_value: &str, - agg_type_key: &str, + agg_type_value: AggregationType, + agg_type_key: AggregationType, agg_id_value: u64, agg_id_key: u64, keys_query: Option, @@ -64,8 +64,8 @@ fn build_context( agg_info: AggregationIdInfo { aggregation_id_for_key: agg_id_key, aggregation_id_for_value: agg_id_value, - aggregation_type_for_key: agg_type_key.to_string(), - aggregation_type_for_value: agg_type_value.to_string(), + aggregation_type_for_key: agg_type_key, + aggregation_type_for_value: agg_type_value, }, do_merge, spatial_filter: String::new(), @@ -235,11 +235,11 @@ fn build_all_test_cases() -> Vec { context: build_context( metric, Statistic::Sum, - vec!["host"], // query_output_labels - vec!["host"], // grouping_labels (store GROUP BY) - vec![], // aggregated_labels (none) - "SumAccumulator", // value accumulator - "SumAccumulator", // key accumulator (same = single) + vec!["host"], // query_output_labels + vec!["host"], // grouping_labels (store GROUP BY) + vec![], // aggregated_labels (none) + AggregationType::Sum, // value accumulator + AggregationType::Sum, // key accumulator (same = single) 42, 42, // same agg_id None, // no keys_query @@ -264,8 +264,8 @@ fn build_all_test_cases() -> Vec { vec!["host"], // query_output_labels vec![], // grouping_labels (no store grouping) vec!["host"], // aggregated_labels (host tracked internally) - "MultipleSumAccumulator", - "MultipleSumAccumulator", // same type = single agg_id + AggregationType::MultipleSum, + AggregationType::MultipleSum, // same type = single agg_id 42, 42, None, @@ -289,8 +289,8 @@ fn build_all_test_cases() -> Vec { vec!["host"], vec![], // grouping_labels (no store grouping) vec!["host"], // aggregated_labels - "CountMinSketch", - "DeltaSetAggregator", + AggregationType::CountMinSketch, + AggregationType::DeltaSetAggregator, 42, 99, // different agg_ids Some(make_keys_query(metric, 99)), @@ -320,8 +320,8 @@ fn build_all_test_cases() -> Vec { vec!["host"], vec!["host"], vec![], - "KLL", - "KLL", + AggregationType::DatasketchesKLL, + AggregationType::DatasketchesKLL, 42, 42, None, @@ -344,8 +344,8 @@ fn build_all_test_cases() -> Vec { vec!["host"], vec![], // no store grouping vec!["host"], // host tracked internally - "HydraKLL", - "DeltaSetAggregator", + AggregationType::HydraKLL, + AggregationType::DeltaSetAggregator, 42, 99, Some(make_keys_query(metric, 99)), @@ -373,8 +373,8 @@ fn build_all_test_cases() -> Vec { vec!["host", "service", "region"], vec!["host", "service", "region"], vec![], - "SumAccumulator", - "SumAccumulator", + AggregationType::Sum, + AggregationType::Sum, 42, 42, None, @@ -397,8 +397,8 @@ fn build_all_test_cases() -> Vec { vec!["host", "service", "region"], vec!["host"], // store groups by host only vec!["service", "region"], // rest tracked internally - "MultipleSumAccumulator", - "MultipleSumAccumulator", + AggregationType::MultipleSum, + AggregationType::MultipleSum, 42, 42, None, @@ -422,8 +422,8 @@ fn build_all_test_cases() -> Vec { vec!["host", "service", "region"], vec!["host"], vec!["service", "region"], - "CountMinSketch", - "DeltaSetAggregator", + AggregationType::CountMinSketch, + AggregationType::DeltaSetAggregator, 42, 99, Some(make_keys_query(metric, 99)), @@ -451,8 +451,8 @@ fn build_all_test_cases() -> Vec { vec!["host", "service", "region"], vec!["host", "service", "region"], vec![], - "KLL", - "KLL", + AggregationType::DatasketchesKLL, + AggregationType::DatasketchesKLL, 42, 42, None, @@ -477,8 +477,8 @@ fn build_all_test_cases() -> Vec { vec!["host", "service", "region"], vec!["host"], vec!["service", "region"], - "HydraKLL", - "DeltaSetAggregator", + AggregationType::HydraKLL, + AggregationType::DeltaSetAggregator, 42, 99, Some(make_keys_query(metric, 99)), diff --git a/asap-query-engine/src/bin/test_e2e_precompute.rs b/asap-query-engine/src/bin/test_e2e_precompute.rs index 3ef70fd..de5006d 100644 --- a/asap-query-engine/src/bin/test_e2e_precompute.rs +++ b/asap-query-engine/src/bin/test_e2e_precompute.rs @@ -9,7 +9,7 @@ //! cargo run --bin test_e2e_precompute use asap_types::aggregation_config::AggregationConfig; -use asap_types::enums::WindowType; +use asap_types::enums::{AggregationType, WindowType}; use prost::Message; use query_engine_rust::data_model::{LockStrategy, QueryLanguage, StreamingConfig}; use query_engine_rust::drivers::ingest::prometheus_remote_write::{ @@ -591,7 +591,7 @@ fn make_sum_agg_config( }; AggregationConfig::new( agg_id, - "SingleSubpopulation".to_string(), + AggregationType::SingleSubpopulation, "Sum".to_string(), HashMap::new(), promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), diff --git a/asap-query-engine/src/bin/test_offline_precomputes.rs b/asap-query-engine/src/bin/test_offline_precomputes.rs index fa25822..dbd2b81 100644 --- a/asap-query-engine/src/bin/test_offline_precomputes.rs +++ b/asap-query-engine/src/bin/test_offline_precomputes.rs @@ -9,7 +9,7 @@ use clap::Parser; use serde::Deserialize; // Internal imports from QueryEngineRust -use promql_utilities::query_logics::enums::{QueryPatternType, Statistic}; +use promql_utilities::query_logics::enums::{AggregationType, QueryPatternType, Statistic}; use query_engine_rust::data_model::{AggregateCore, KeyByLabelValues, PrecomputedOutput}; use query_engine_rust::precompute_operators::*; @@ -740,7 +740,7 @@ fn query_precompute_for_statistic( query_kwargs: &HashMap, ) -> Result> { match precompute.get_accumulator_type() { - "SumAccumulator" => { + AggregationType::Sum => { let acc = precompute .as_any() .downcast_ref::() @@ -749,7 +749,7 @@ fn query_precompute_for_statistic( acc.query(*statistic, None) .map_err(|e| format!("{}", e).into()) } - "MultipleIncreaseAccumulator" => { + AggregationType::MultipleIncrease => { let acc = precompute .as_any() .downcast_ref::() @@ -761,7 +761,7 @@ fn query_precompute_for_statistic( acc.query(*statistic, key_val, Some(query_kwargs)) .map_err(|e| format!("{}", e).into()) } - "CountMinSketchAccumulator" => { + AggregationType::CountMinSketch => { let acc = precompute .as_any() .downcast_ref::() @@ -773,7 +773,7 @@ fn query_precompute_for_statistic( acc.query(*statistic, key_val, Some(query_kwargs)) .map_err(|e| format!("{}", e).into()) } - "CountMinSketchWithHeapAccumulator" => { + AggregationType::CountMinSketchWithHeap => { let acc = precompute .as_any() .downcast_ref::() @@ -785,7 +785,7 @@ fn query_precompute_for_statistic( acc.query(*statistic, key_val, Some(query_kwargs)) .map_err(|e| format!("{}", e).into()) } - "DatasketchesKLLAccumulator" => { + AggregationType::DatasketchesKLL => { let acc = precompute .as_any() .downcast_ref::() @@ -794,7 +794,7 @@ fn query_precompute_for_statistic( acc.query(*statistic, Some(query_kwargs)) .map_err(|e| format!("{}", e).into()) } - "DeltaSetAggregatorAccumulator" => { + AggregationType::DeltaSetAggregator => { let acc = precompute .as_any() .downcast_ref::() @@ -807,7 +807,7 @@ fn query_precompute_for_statistic( Ok((acc.added.union(&acc.removed).count()) as f64) } } - "SetAggregatorAccumulator" => { + AggregationType::SetAggregator => { let acc = precompute .as_any() .downcast_ref::() @@ -820,11 +820,7 @@ fn query_precompute_for_statistic( Ok(acc.added.len() as f64) } } - _ => Err(format!( - "Unsupported accumulator type: {}", - precompute.get_accumulator_type() - ) - .into()), + other => Err(format!("Unsupported accumulator type: {other:?}").into()), } } diff --git a/asap-query-engine/src/data_model/enums.rs b/asap-query-engine/src/data_model/enums.rs index 3cc571f..885fbee 100644 --- a/asap-query-engine/src/data_model/enums.rs +++ b/asap-query-engine/src/data_model/enums.rs @@ -10,6 +10,7 @@ pub enum StreamingEngine { } pub use asap_types::enums::{CleanupPolicy, QueryLanguage, WindowType}; +pub use promql_utilities::query_logics::enums::AggregationType; #[derive(clap::ValueEnum, Clone, Debug, PartialEq)] pub enum QueryProtocol { diff --git a/asap-query-engine/src/data_model/precomputed_output.rs b/asap-query-engine/src/data_model/precomputed_output.rs index d57c67f..88810da 100644 --- a/asap-query-engine/src/data_model/precomputed_output.rs +++ b/asap-query-engine/src/data_model/precomputed_output.rs @@ -5,7 +5,7 @@ use std::io::Read as _; use tracing::error; use crate::data_model::traits::SerializableToSink; -use crate::data_model::{KeyByLabelValues, StreamingConfig}; +use crate::data_model::{AggregationType, KeyByLabelValues, StreamingConfig}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PrecomputedOutput { @@ -231,7 +231,7 @@ impl PrecomputedOutput { .map_err(|e| format!("Failed to decompress precompute data: {e}"))?; let precompute = Self::create_precompute_from_bytes( - &config.aggregation_type, + config.aggregation_type, Vec::as_slice(&precompute_bytes), )?; @@ -362,34 +362,34 @@ impl PrecomputedOutput { /// Factory method to create precompute accumulator from bytes fn create_precompute_from_bytes( - precompute_type: &str, + precompute_type: AggregationType, buffer: &[u8], ) -> Result, Box> { use crate::precompute_operators::*; match precompute_type { - "Sum" | "sum" => { + AggregationType::Sum => { let accumulator = SumAccumulator::deserialize_from_bytes_arroyo(buffer) .map_err(|e| format!("Failed to deserialize SumAccumulator: {e}"))?; Ok(Box::new(accumulator)) } - "MinMax" => { + AggregationType::MinMax => { let accumulator = MinMaxAccumulator::deserialize_from_bytes(buffer) .map_err(|e| format!("Failed to deserialize MinMaxAccumulator: {e}"))?; Ok(Box::new(accumulator)) } - "Increase" => { + AggregationType::Increase => { let accumulator = IncreaseAccumulator::deserialize_from_bytes(buffer) .map_err(|e| format!("Failed to deserialize IncreaseAccumulator: {e}"))?; Ok(Box::new(accumulator)) } - "MultipleSum" => { + AggregationType::MultipleSum => { let accumulator = MultipleSumAccumulator::deserialize_from_bytes_arroyo(buffer) .map_err(|e| format!("Failed to deserialize MultipleSumAccumulator: {e}"))?; Ok(Box::new(accumulator)) } - "MultipleMinMax" => { + AggregationType::MultipleMinMax => { let accumulator = MultipleMinMaxAccumulator::deserialize_from_bytes(buffer, "min".to_string()) .map_err(|e| { @@ -397,19 +397,19 @@ impl PrecomputedOutput { })?; Ok(Box::new(accumulator)) } - "MultipleIncrease" => { + AggregationType::MultipleIncrease => { let accumulator = MultipleIncreaseAccumulator::deserialize_from_bytes_arroyo( buffer, ) .map_err(|e| format!("Failed to deserialize MultipleIncreaseAccumulator: {e}"))?; Ok(Box::new(accumulator)) } - "CountMinSketch" => { + AggregationType::CountMinSketch => { let accumulator = CountMinSketchAccumulator::deserialize_from_bytes_arroyo(buffer) .map_err(|e| format!("Failed to deserialize CountMinSketchAccumulator: {e}"))?; Ok(Box::new(accumulator)) } - "CountMinSketchWithHeap" => { + AggregationType::CountMinSketchWithHeap => { let accumulator = CountMinSketchWithHeapAccumulator::deserialize_from_bytes_arroyo(buffer) .map_err(|e| { @@ -417,26 +417,26 @@ impl PrecomputedOutput { })?; Ok(Box::new(accumulator)) } - "DatasketchesKLL" => { + AggregationType::DatasketchesKLL => { let accumulator = DatasketchesKLLAccumulator::deserialize_from_bytes_arroyo(buffer) .map_err(|e| { format!("Failed to deserialize DatasketchesKLLAccumulator: {e}") })?; Ok(Box::new(accumulator)) } - "HydraKLL" => { + AggregationType::HydraKLL => { let accumulator = HydraKllSketchAccumulator::deserialize_from_bytes_arroyo(buffer) .map_err(|e| format!("Failed to deserialize HydraKllSketchAccumulator: {e}"))?; Ok(Box::new(accumulator)) } - "DeltaSetAggregator" => { + AggregationType::DeltaSetAggregator => { let accumulator = DeltaSetAggregatorAccumulator::deserialize_from_bytes_arroyo( buffer, ) .map_err(|e| format!("Failed to deserialize DeltaSetAggregatorAccumulator: {e}"))?; Ok(Box::new(accumulator)) } - _ => Err(format!("Unknown precompute type: {precompute_type}").into()), + _ => Err(format!("Unknown precompute type: {precompute_type:?}").into()), } } } diff --git a/asap-query-engine/src/data_model/traits.rs b/asap-query-engine/src/data_model/traits.rs index 2620536..c47fcef 100644 --- a/asap-query-engine/src/data_model/traits.rs +++ b/asap-query-engine/src/data_model/traits.rs @@ -2,7 +2,7 @@ use crate::data_model::KeyByLabelValues; use serde_json::Value; use std::collections::HashMap; -use promql_utilities::query_logics::enums::Statistic; +use promql_utilities::query_logics::enums::{AggregationType, Statistic}; pub use asap_types::traits::SerializableToSink; @@ -26,7 +26,7 @@ pub trait AggregateCore: SerializableToSink + Send + Sync { ) -> Result, Box>; /// Get the accumulator type identifier for merge compatibility checking - fn get_accumulator_type(&self) -> &'static str; + fn get_accumulator_type(&self) -> AggregationType; /// Get all keys stored in this accumulator fn get_keys(&self) -> Option>; diff --git a/asap-query-engine/src/engines/logical/plan_builder.rs b/asap-query-engine/src/engines/logical/plan_builder.rs index cdf6be0..813f024 100644 --- a/asap-query-engine/src/engines/logical/plan_builder.rs +++ b/asap-query-engine/src/engines/logical/plan_builder.rs @@ -15,7 +15,7 @@ use datafusion_summary_library::{ InferOperation, PrecomputedSummaryRead, SketchType, SummaryInfer, SummaryMergeMultiple, }; use promql_parser::parser::token::{self, T_ADD, T_DIV, T_MOD, T_MUL, T_POW, T_SUB}; -use promql_utilities::query_logics::enums::Statistic; +use promql_utilities::query_logics::enums::{AggregationType, Statistic}; use std::sync::Arc; use crate::engines::simple_engine::{QueryExecutionContext, StoreQueryParams}; @@ -215,33 +215,34 @@ impl QueryExecutionContext { } } - /// Map aggregation type string to SketchType (SummaryType) for the value accumulator + /// Map aggregation type to SketchType (SummaryType) for the value accumulator fn map_aggregation_type_to_summary_type(&self) -> Result { - Self::agg_type_str_to_sketch_type(&self.agg_info.aggregation_type_for_value) + Self::agg_type_to_sketch_type(self.agg_info.aggregation_type_for_value) } - /// Map aggregation type string to SketchType (SummaryType) for the key accumulator + /// Map aggregation type to SketchType (SummaryType) for the key accumulator fn map_key_aggregation_type_to_summary_type(&self) -> Result { - Self::agg_type_str_to_sketch_type(&self.agg_info.aggregation_type_for_key) + Self::agg_type_to_sketch_type(self.agg_info.aggregation_type_for_key) } - fn agg_type_str_to_sketch_type(agg_type: &str) -> Result { + fn agg_type_to_sketch_type(agg_type: AggregationType) -> Result { match agg_type { - "SumAggregator" | "SumAccumulator" | "Sum" => Ok(SketchType::Sum), - "IncreaseAggregator" | "IncreaseAccumulator" | "Increase" => Ok(SketchType::Increase), - "MinMaxAggregator" | "MinMaxAccumulator" | "MinMax" => Ok(SketchType::MinMax), - "MultipleSumAccumulator" | "MultipleSum" => Ok(SketchType::MultipleSum), - "MultipleIncreaseAccumulator" | "MultipleIncrease" => Ok(SketchType::MultipleIncrease), - "MultipleMinMaxAccumulator" | "MultipleMinMax" => Ok(SketchType::MultipleMinMax), - "DeltaSetAggregator" => Ok(SketchType::DeltaSetAggregator), - "SetAggregator" => Ok(SketchType::SetAggregator), - "DatasketchesKLLAccumulator" | "DatasketchesKLL" | "KLL" => Ok(SketchType::KLL), - "HydraKllSketchAccumulator" | "HydraKLL" => Ok(SketchType::HydraKLL), - "CountMinSketchAccumulator" | "CountMinSketch" => Ok(SketchType::CountMinSketch), - "HyperLogLog" | "HLL" => Ok(SketchType::HLL), + AggregationType::Sum => Ok(SketchType::Sum), + AggregationType::Increase => Ok(SketchType::Increase), + AggregationType::MinMax => Ok(SketchType::MinMax), + AggregationType::MultipleSum => Ok(SketchType::MultipleSum), + AggregationType::MultipleIncrease => Ok(SketchType::MultipleIncrease), + AggregationType::MultipleMinMax => Ok(SketchType::MultipleMinMax), + AggregationType::DeltaSetAggregator => Ok(SketchType::DeltaSetAggregator), + AggregationType::SetAggregator => Ok(SketchType::SetAggregator), + AggregationType::DatasketchesKLL => Ok(SketchType::KLL), + AggregationType::HydraKLL => Ok(SketchType::HydraKLL), + AggregationType::CountMinSketch | AggregationType::CountMinSketchWithHeap => { + Ok(SketchType::CountMinSketch) + } + AggregationType::HLL => Ok(SketchType::HLL), _ => Err(DataFusionError::Plan(format!( - "Unknown aggregation type: {}", - agg_type + "Unknown aggregation type: {agg_type:?}" ))), } } @@ -363,16 +364,23 @@ mod tests { metric: &str, statistic: Statistic, output_labels: Vec<&str>, - aggregation_type: &str, + aggregation_type: AggregationType, ) -> QueryExecutionContext { - create_test_context_with_keys(metric, statistic, output_labels, aggregation_type, None, "") + create_test_context_with_keys( + metric, + statistic, + output_labels, + aggregation_type, + None, + aggregation_type, + ) } fn create_test_context_with_kwargs( metric: &str, statistic: Statistic, output_labels: Vec<&str>, - aggregation_type: &str, + aggregation_type: AggregationType, kwargs: HashMap, ) -> QueryExecutionContext { let mut ctx = create_test_context_with_keys( @@ -381,7 +389,7 @@ mod tests { output_labels, aggregation_type, None, - "", + aggregation_type, ); ctx.metadata.query_kwargs = kwargs; ctx @@ -391,9 +399,9 @@ mod tests { metric: &str, statistic: Statistic, output_labels: Vec<&str>, - aggregation_type_for_value: &str, + aggregation_type_for_value: AggregationType, keys_query: Option, - aggregation_type_for_key: &str, + aggregation_type_for_key: AggregationType, ) -> QueryExecutionContext { // Default: grouping_labels == query_output_labels create_test_context_with_keys_and_grouping( @@ -411,9 +419,9 @@ mod tests { metric: &str, statistic: Statistic, output_labels: Vec<&str>, - aggregation_type_for_value: &str, + aggregation_type_for_value: AggregationType, keys_query: Option, - aggregation_type_for_key: &str, + aggregation_type_for_key: AggregationType, grouping_label_strs: Vec<&str>, ) -> QueryExecutionContext { // Default: aggregated_labels = output_labels - grouping_labels @@ -439,9 +447,9 @@ mod tests { metric: &str, statistic: Statistic, output_labels: Vec<&str>, - aggregation_type_for_value: &str, + aggregation_type_for_value: AggregationType, keys_query: Option, - aggregation_type_for_key: &str, + aggregation_type_for_key: AggregationType, grouping_label_strs: Vec<&str>, aggregated_label_strs: Vec<&str>, ) -> QueryExecutionContext { @@ -482,8 +490,8 @@ mod tests { agg_info: AggregationIdInfo { aggregation_id_for_key, aggregation_id_for_value: 42, - aggregation_type_for_key: aggregation_type_for_key.to_string(), - aggregation_type_for_value: aggregation_type_for_value.to_string(), + aggregation_type_for_key, + aggregation_type_for_value, }, do_merge: false, spatial_filter: String::new(), @@ -499,7 +507,7 @@ mod tests { "http_requests", Statistic::Sum, vec!["host"], - "SumAggregator", + AggregationType::Sum, ); let plan = context.to_logical_plan().unwrap(); @@ -519,7 +527,7 @@ mod tests { "http_requests", Statistic::Sum, vec!["host", "region", "service"], - "SumAggregator", + AggregationType::Sum, ); let plan = context.to_logical_plan().unwrap(); @@ -530,19 +538,30 @@ mod tests { #[test] fn test_map_statistic_to_infer_operation() { - let context = create_test_context("test", Statistic::Sum, vec!["host"], "SumAggregator"); + let context = + create_test_context("test", Statistic::Sum, vec!["host"], AggregationType::Sum); assert!(matches!( context.map_statistic_to_infer_operation().unwrap(), InferOperation::ExtractSum )); - let context = create_test_context("test", Statistic::Min, vec!["host"], "MinMaxAggregator"); + let context = create_test_context( + "test", + Statistic::Min, + vec!["host"], + AggregationType::MinMax, + ); assert!(matches!( context.map_statistic_to_infer_operation().unwrap(), InferOperation::ExtractMin )); - let context = create_test_context("test", Statistic::Max, vec!["host"], "MinMaxAggregator"); + let context = create_test_context( + "test", + Statistic::Max, + vec!["host"], + AggregationType::MinMax, + ); assert!(matches!( context.map_statistic_to_infer_operation().unwrap(), InferOperation::ExtractMax @@ -551,7 +570,8 @@ mod tests { #[test] fn test_map_aggregation_type_to_summary_type() { - let context = create_test_context("test", Statistic::Sum, vec!["host"], "SumAggregator"); + let context = + create_test_context("test", Statistic::Sum, vec!["host"], AggregationType::Sum); assert_eq!( context.map_aggregation_type_to_summary_type().unwrap(), SketchType::Sum @@ -561,7 +581,7 @@ mod tests { "test", Statistic::Increase, vec!["host"], - "IncreaseAggregator", + AggregationType::Increase, ); assert_eq!( context.map_aggregation_type_to_summary_type().unwrap(), @@ -572,7 +592,7 @@ mod tests { "test", Statistic::Quantile, vec!["host"], - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, ); assert_eq!( context.map_aggregation_type_to_summary_type().unwrap(), @@ -673,7 +693,7 @@ mod tests { "http_requests", Statistic::Sum, vec!["host"], - "MultipleSumAccumulator", + AggregationType::MultipleSum, ); assert_eq!( context.map_aggregation_type_to_summary_type().unwrap(), @@ -690,7 +710,7 @@ mod tests { "http_requests", Statistic::Sum, vec!["host"], - "MultipleSumAccumulator", + AggregationType::MultipleSum, ); let plan = context.to_logical_plan().unwrap(); @@ -722,7 +742,7 @@ mod tests { "http_requests", Statistic::Sum, vec!["host"], - "MultipleSumAccumulator", + AggregationType::MultipleSum, ); let plan = context.to_logical_plan().unwrap(); @@ -748,7 +768,7 @@ mod tests { "http_requests", Statistic::Count, vec!["host"], - "CountMinSketchAccumulator", + AggregationType::CountMinSketch, ); assert_eq!( context.map_aggregation_type_to_summary_type().unwrap(), @@ -764,7 +784,7 @@ mod tests { "http_requests", Statistic::Count, vec!["host"], - "CountMinSketchAccumulator", + AggregationType::CountMinSketch, ); let plan = context.to_logical_plan().unwrap(); @@ -791,7 +811,7 @@ mod tests { "http_requests", Statistic::Count, vec!["host"], - "CountMinSketchAccumulator", + AggregationType::CountMinSketch, ); let plan = context.to_logical_plan().unwrap(); @@ -821,9 +841,9 @@ mod tests { "http_requests", Statistic::Sum, vec!["host"], - "MultipleSumAccumulator", // values use Hydra + AggregationType::MultipleSum, // values use Hydra Some(keys_query), - "DeltaSetAggregator", // keys use DeltaSet + AggregationType::DeltaSetAggregator, // keys use DeltaSet ); let plan = context.to_logical_plan().unwrap(); @@ -876,9 +896,9 @@ mod tests { "request_duration", Statistic::Quantile, vec!["host", "endpoint"], // query_output_labels - "HydraKllSketchAccumulator", + AggregationType::HydraKLL, Some(keys_query), - "DeltaSetAggregator", + AggregationType::DeltaSetAggregator, vec!["host"], // grouping_labels (spatial) ); context.metadata.query_kwargs = kwargs; @@ -928,9 +948,9 @@ mod tests { "request_duration", Statistic::Quantile, vec!["host", "endpoint"], // query_output_labels - "HydraKllSketchAccumulator", + AggregationType::HydraKLL, Some(keys_query), - "DeltaSetAggregator", + AggregationType::DeltaSetAggregator, vec!["host"], // grouping_labels (spatial store labels) ); @@ -961,7 +981,7 @@ mod tests { "latency", Statistic::Quantile, vec!["host"], - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, kwargs, ); @@ -980,7 +1000,7 @@ mod tests { "latency", Statistic::Quantile, vec!["host"], - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, ); match context.map_statistic_to_infer_operation().unwrap() { diff --git a/asap-query-engine/src/engines/physical/accumulator_serde.rs b/asap-query-engine/src/engines/physical/accumulator_serde.rs index 6955cbd..cc7f85e 100644 --- a/asap-query-engine/src/engines/physical/accumulator_serde.rs +++ b/asap-query-engine/src/engines/physical/accumulator_serde.rs @@ -279,6 +279,7 @@ pub fn deserialize_keys_accumulator( #[cfg(test)] mod tests { use super::*; + use crate::data_model::AggregationType; // Helper to serialize f64 as MessagePack (Arroyo format) fn serialize_f64_arroyo(value: f64) -> Vec { @@ -292,7 +293,7 @@ mod tests { let restored = deserialize_accumulator(&bytes, &SketchType::Sum).unwrap(); - assert_eq!(restored.get_accumulator_type(), "SumAccumulator"); + assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); } #[test] diff --git a/asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs b/asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs index a58c984..f335482 100644 --- a/asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs +++ b/asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs @@ -287,7 +287,7 @@ pub(crate) fn build_output_batch( #[cfg(test)] mod tests { use super::*; - use crate::data_model::KeyByLabelValues; + use crate::data_model::{AggregationType, KeyByLabelValues}; use crate::engines::physical::accumulator_serde::serialize_accumulator_arroyo; use crate::precompute_operators::{ DatasketchesKLLAccumulator, SetAggregatorAccumulator, SumAccumulator, @@ -365,7 +365,7 @@ mod tests { // The single row should pass through unchanged let (_, merged_bytes) = groups.values().next().unwrap(); let restored = deserialize_accumulator(merged_bytes, &SketchType::Sum).unwrap(); - assert_eq!(restored.get_accumulator_type(), "SumAccumulator"); + assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); } #[test] diff --git a/asap-query-engine/src/engines/simple_engine.rs b/asap-query-engine/src/engines/simple_engine.rs index 83e23d8..64da6eb 100644 --- a/asap-query-engine/src/engines/simple_engine.rs +++ b/asap-query-engine/src/engines/simple_engine.rs @@ -9,7 +9,7 @@ use crate::engines::query_result::{InstantVectorElement, QueryResult, RangeVecto use crate::stores::{Store, TimestampedBucketsMap}; use core::panic; use promql_utilities::get_is_collapsable; -use promql_utilities::query_logics::enums::{AggregationOperator, PromQLFunction}; +use promql_utilities::query_logics::enums::{AggregationOperator, AggregationType, PromQLFunction}; use serde_json::Value; use std::collections::HashMap; use std::sync::Arc; @@ -563,12 +563,12 @@ impl SimpleEngine { end_timestamp: u64, agg_info: &AggregationIdInfo, ) -> Result { - let (start_timestamp, end_timestamp) = match agg_info.aggregation_type_for_key.as_str() { - "DeltaSetAggregator" => { + let (start_timestamp, end_timestamp) = match agg_info.aggregation_type_for_key { + AggregationType::DeltaSetAggregator => { // All keys from beginning of time (0, end_timestamp) } - "SetAggregator" => { + AggregationType::SetAggregator => { // Latest window only let window_size = self .streaming_config @@ -583,7 +583,7 @@ impl SimpleEngine { (end_timestamp - window_size, end_timestamp) } other => { - return Err(format!("Unsupported key aggregation type: {}", other)); + return Err(format!("Unsupported key aggregation type: {other:?}")); } }; @@ -770,7 +770,7 @@ impl SimpleEngine { self.merge_precomputed_outputs( &values_map, do_merge, - agg_info.aggregation_type_for_value.clone(), + agg_info.aggregation_type_for_value, ) }; @@ -805,7 +805,7 @@ impl SimpleEngine { let merged = self.merge_precomputed_outputs( &keys_map, do_merge, - agg_info.aggregation_type_for_key.clone(), + agg_info.aggregation_type_for_key, ); debug!( "[LATENCY] Keys merge operation: {:.2}ms, resulted in {} merged outputs", @@ -1645,8 +1645,8 @@ impl SimpleEngine { let query_config_aggregations = &query_config.aggregations; let mut aggregation_id_for_key: Option = None; let mut aggregation_id_for_value: Option = None; - let mut aggregation_type_for_key: Option = None; - let mut aggregation_type_for_value: Option = None; + let mut aggregation_type_for_key: Option = None; + let mut aggregation_type_for_value: Option = None; if query_config_aggregations.is_empty() { panic!("Query config for query has no aggregations defined",); @@ -1657,11 +1657,12 @@ impl SimpleEngine { let aggregation_type = self .streaming_config .get_aggregation_config(aggregation.aggregation_id) - .map(|config| config.aggregation_type.clone()); + .map(|config| config.aggregation_type); - if aggregation_type.as_ref().unwrap() == "DeltaSetAggregator" - || aggregation_type.as_ref().unwrap() == "SetAggregator" - { + if matches!( + aggregation_type.as_ref().unwrap(), + AggregationType::DeltaSetAggregator | AggregationType::SetAggregator + ) { if aggregation_id_for_key.is_some() { panic!("Aggregation ID for key must be None"); } @@ -1685,11 +1686,11 @@ impl SimpleEngine { aggregation_type_for_key = self .streaming_config .get_aggregation_config(aggregation_id_for_key.unwrap()) - .map(|config| config.aggregation_type.clone()); + .map(|config| config.aggregation_type); aggregation_type_for_value = self .streaming_config .get_aggregation_config(aggregation_id_for_value.unwrap()) - .map(|config| config.aggregation_type.clone()); + .map(|config| config.aggregation_type); } // check for None @@ -2871,7 +2872,7 @@ impl SimpleEngine { &self, precomputed_outputs_map: &TimestampedBucketsMap, do_merge: bool, - aggregation_type: String, + aggregation_type: AggregationType, ) -> HashMap, Box> { #[cfg(feature = "extra_debugging")] let start_time = Instant::now(); @@ -2879,29 +2880,16 @@ impl SimpleEngine { debug!("Starting merge for {} keys", precomputed_outputs_map.len()); #[cfg(feature = "extra_debugging")] debug!( - "do_merge: {}, aggregation_type: {}", + "do_merge: {}, aggregation_type: {:?}", do_merge, aggregation_type ); // Merge if: temporal query OR DeltaSetAggregator (which accumulates keys over time) - let should_merge = do_merge || aggregation_type == "DeltaSetAggregator"; + let should_merge = do_merge || aggregation_type == AggregationType::DeltaSetAggregator; let mut merged = HashMap::with_capacity(precomputed_outputs_map.len()); - for (idx, (key, timestamped_buckets)) in precomputed_outputs_map.iter().enumerate() { - #[cfg(feature = "extra_debugging")] - debug!( - "Processing key {} of {}: {:?}", - idx + 1, - precomputed_outputs_map.len(), - key - ); - #[cfg(feature = "extra_debugging")] - debug!( - " Number of precomputes for this key: {}", - timestamped_buckets.len() - ); - + for (key, timestamped_buckets) in precomputed_outputs_map.iter() { if !timestamped_buckets.is_empty() { // Extract just the buckets (without timestamps) for merging let precomputes: Vec> = timestamped_buckets @@ -2963,7 +2951,7 @@ impl SimpleEngine { // Try to use optimized batch merge for KLL accumulators if !accumulators.is_empty() - && accumulators[0].get_accumulator_type() == "DatasketchesKLLAccumulator" + && accumulators[0].get_accumulator_type() == AggregationType::DatasketchesKLL { use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; @@ -2981,7 +2969,7 @@ impl SimpleEngine { // Try to use optimized batch merge for CountMinSketch accumulators if !accumulators.is_empty() - && accumulators[0].get_accumulator_type() == "CountMinSketchAccumulator" + && accumulators[0].get_accumulator_type() == AggregationType::CountMinSketch { use crate::precompute_operators::count_min_sketch_accumulator::CountMinSketchAccumulator; @@ -3133,7 +3121,7 @@ impl SimpleEngine { // Handle different accumulator types and statistics using the trait methods // TODO: change this logic to just check Single vs MultipleSubpopulationAggregate match precompute.get_accumulator_type() { - "SumAccumulator" => { + AggregationType::Sum => { if let Some(sum_acc) = precompute.as_any().downcast_ref::() { use crate::data_model::SingleSubpopulationAggregate; sum_acc.query(*statistic, None) @@ -3141,7 +3129,7 @@ impl SimpleEngine { Err("Failed to downcast to SumAccumulator".into()) } } - "MinMaxAccumulator" => { + AggregationType::MinMax => { if let Some(minmax_acc) = precompute.as_any().downcast_ref::() { use crate::data_model::SingleSubpopulationAggregate; minmax_acc.query(*statistic, None) @@ -3149,7 +3137,7 @@ impl SimpleEngine { Err("Failed to downcast to MinMaxAccumulator".into()) } } - "IncreaseAccumulator" => { + AggregationType::Increase => { if let Some(inc_acc) = precompute.as_any().downcast_ref::() { use crate::data_model::SingleSubpopulationAggregate; inc_acc.query(*statistic, None) @@ -3157,7 +3145,7 @@ impl SimpleEngine { Err("Failed to downcast to IncreaseAccumulator".into()) } } - "MultipleSumAccumulator" => { + AggregationType::MultipleSum => { if let Some(multi_sum_acc) = precompute.as_any().downcast_ref::() { if let Some(key_val) = key { use crate::data_model::MultipleSubpopulationAggregate; @@ -3169,7 +3157,7 @@ impl SimpleEngine { Err("Failed to downcast to MultipleSumAccumulator".into()) } } - "MultipleMinMaxAccumulator" => { + AggregationType::MultipleMinMax => { if let Some(multi_minmax_acc) = precompute.as_any().downcast_ref::() { if let Some(key_val) = key { use crate::data_model::MultipleSubpopulationAggregate; @@ -3181,7 +3169,7 @@ impl SimpleEngine { Err("Failed to downcast to MultipleMinMaxAccumulator".into()) } } - "MultipleIncreaseAccumulator" => { + AggregationType::MultipleIncrease => { if let Some(multi_inc_acc) = precompute.as_any().downcast_ref::() { if let Some(key_val) = key { use crate::data_model::MultipleSubpopulationAggregate; @@ -3193,7 +3181,7 @@ impl SimpleEngine { Err("Failed to downcast to MultipleIncreaseAccumulator".into()) } } - "CountMinSketchAccumulator" => { + AggregationType::CountMinSketch => { if let Some(cms_acc) = precompute.as_any().downcast_ref::() { use crate::data_model::MultipleSubpopulationAggregate; if let Some(key_val) = key { @@ -3205,7 +3193,7 @@ impl SimpleEngine { Err("Failed to downcast to CountMinSketchAccumulator".into()) } } - "CountMinSketchWithHeapAccumulator" => { + AggregationType::CountMinSketchWithHeap => { if let Some(cms_heap_acc) = precompute.as_any().downcast_ref::() { use crate::data_model::MultipleSubpopulationAggregate; if let Some(key_val) = key { @@ -3217,7 +3205,7 @@ impl SimpleEngine { Err("Failed to downcast to CountMinSketchWithHeapAccumulator".into()) } } - "DatasketchesKLLAccumulator" => { + AggregationType::DatasketchesKLL => { if let Some(kll_acc) = precompute.as_any().downcast_ref::() { use crate::data_model::SingleSubpopulationAggregate; kll_acc.query(*statistic, Some(query_kwargs)) @@ -3225,7 +3213,7 @@ impl SimpleEngine { Err("Failed to downcast to DatasketchesKLLAccumulator".into()) } } - "HydraKllSketchAccumulator" => { + AggregationType::HydraKLL => { if let Some(hydra_kll_acc) = precompute.as_any() .downcast_ref::() { @@ -3239,7 +3227,7 @@ impl SimpleEngine { Err("Failed to downcast to HydraKllSketchAccumulator".into()) } } - "DeltaSetAggregatorAccumulator" => { + AggregationType::DeltaSetAggregator => { if let Some(delta_acc) = precompute.as_any().downcast_ref::() { if let Some(key_val) = key { use crate::data_model::MultipleSubpopulationAggregate; @@ -3252,7 +3240,7 @@ impl SimpleEngine { Err("Failed to downcast to DeltaSetAggregatorAccumulator".into()) } } - "SetAggregatorAccumulator" => { + AggregationType::SetAggregator => { if let Some(set_acc) = precompute.as_any().downcast_ref::() { if let Some(key_val) = key { use crate::data_model::MultipleSubpopulationAggregate; @@ -3265,8 +3253,8 @@ impl SimpleEngine { Err("Failed to downcast to SetAggregatorAccumulator".into()) } } - _ => { - Err(format!("Unknown accumulator type: {}", precompute.get_accumulator_type()).into()) + other => { + Err(format!("Unknown accumulator type: {other:?}").into()) } } } @@ -3677,7 +3665,7 @@ impl SimpleEngine { if !window_buckets.is_empty() { // Merge available buckets - let mut merger = create_window_merger(accumulator_type); + let mut merger = create_window_merger(*accumulator_type); merger.initialize(window_buckets); match merger.get_merged() { @@ -3741,7 +3729,7 @@ impl SimpleEngine { #[cfg(test)] mod range_query_tests { - use crate::data_model::{AggregateCore, KeyByLabelValues, SerializableToSink}; + use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; use crate::engines::window_merger::NaiveMerger; use serde_json::Value; use std::any::Any; @@ -3797,8 +3785,8 @@ mod range_query_tests { } } - fn get_accumulator_type(&self) -> &'static str { - "MockBucketAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Sum } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/engines/window_merger.rs b/asap-query-engine/src/engines/window_merger.rs index 976cec8..56e75b2 100644 --- a/asap-query-engine/src/engines/window_merger.rs +++ b/asap-query-engine/src/engines/window_merger.rs @@ -8,7 +8,7 @@ //! - `IncrementalMerger`: Add/subtract for subtractable accumulators (future) //! - `SwagMerger`: Two-stack queue for non-subtractable accumulators (future) -use crate::data_model::AggregateCore; +use crate::data_model::{AggregateCore, AggregationType}; /// Trait for merging buckets in a sliding window /// @@ -94,7 +94,7 @@ impl WindowMerger for NaiveMerger { /// - IncrementalMerger for subtractable accumulators (Sum, CountMinSketch) /// - SwagMerger for non-subtractable accumulators (KLL, MinMax) #[allow(dead_code)] -pub fn create_window_merger(_accumulator_type: &str) -> Box { +pub fn create_window_merger(_accumulator_type: AggregationType) -> Box { // Future implementation: // match accumulator_type { // "SumAccumulator" | "CountMinSketchAccumulator" => Box::new(IncrementalMerger::new()), @@ -159,8 +159,8 @@ mod tests { } } - fn get_accumulator_type(&self) -> &'static str { - "MockSumAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Sum } fn get_keys(&self) -> Option> { @@ -186,7 +186,7 @@ mod tests { #[test] fn test_create_window_merger() { - let merger = create_window_merger("SumAccumulator"); + let merger = create_window_merger(AggregationType::Sum); assert!(!merger.is_initialized()); } diff --git a/asap-query-engine/src/precompute_engine/accumulator_factory.rs b/asap-query-engine/src/precompute_engine/accumulator_factory.rs index 77d2ca7..4852126 100644 --- a/asap-query-engine/src/precompute_engine/accumulator_factory.rs +++ b/asap-query-engine/src/precompute_engine/accumulator_factory.rs @@ -1,4 +1,4 @@ -use crate::data_model::{AggregateCore, KeyByLabelValues, Measurement}; +use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, Measurement}; use crate::precompute_operators::{ CountMinSketchAccumulator, DatasketchesKLLAccumulator, HydraKllSketchAccumulator, IncreaseAccumulator, MinMaxAccumulator, MultipleIncreaseAccumulator, MultipleMinMaxAccumulator, @@ -515,20 +515,14 @@ impl AccumulatorUpdater for HydraKllAccumulatorUpdater { /// in the corresponding struct. pub fn config_is_keyed(config: &AggregationConfig) -> bool { matches!( - config.aggregation_type.as_str(), - "MultipleSubpopulation" - | "MultipleSum" - | "multiple_sum" - | "MultipleIncrease" - | "multiple_increase" - | "MultipleMinMax" - | "multiple_min_max" - | "CountMinSketch" - | "count_min_sketch" - | "CMS" - | "cms" - | "HydraKLL" - | "hydra_kll" + config.aggregation_type, + AggregationType::MultipleSubpopulation + | AggregationType::MultipleSum + | AggregationType::MultipleIncrease + | AggregationType::MultipleMinMax + | AggregationType::CountMinSketch + | AggregationType::CountMinSketchWithHeap + | AggregationType::HydraKLL ) } @@ -570,11 +564,10 @@ fn hydra_kll_params(config: &AggregationConfig) -> (usize, usize, u16) { /// Create an appropriate `AccumulatorUpdater` from an `AggregationConfig`. pub fn create_accumulator_updater(config: &AggregationConfig) -> Box { - let agg_type = config.aggregation_type.as_str(); let sub_type = config.aggregation_sub_type.as_str(); - match agg_type { - "SingleSubpopulation" => match sub_type { + match config.aggregation_type { + AggregationType::SingleSubpopulation => match sub_type { "Sum" | "sum" => Box::new(SumAccumulatorUpdater::new()), "Min" | "min" => Box::new(MinMaxAccumulatorUpdater::new("min".to_string())), "Max" | "max" => Box::new(MinMaxAccumulatorUpdater::new("max".to_string())), @@ -590,7 +583,7 @@ pub fn create_accumulator_updater(config: &AggregationConfig) -> Box match sub_type { + AggregationType::MultipleSubpopulation => match sub_type { "Sum" | "sum" => Box::new(MultipleSumUpdater::new()), "Min" | "min" => Box::new(MultipleMinMaxUpdater::new("min".to_string())), "Max" | "max" => Box::new(MultipleMinMaxUpdater::new("max".to_string())), @@ -611,34 +604,38 @@ pub fn create_accumulator_updater(config: &AggregationConfig) -> Box { + AggregationType::DatasketchesKLL => { Box::new(KllAccumulatorUpdater::new(kll_k_param(config))) } - "MultipleSum" | "multiple_sum" => Box::new(MultipleSumUpdater::new()), - "MultipleIncrease" | "multiple_increase" => Box::new(MultipleIncreaseUpdater::new()), - "MultipleMinMax" | "multiple_min_max" => Box::new(MultipleMinMaxUpdater::new( + AggregationType::MultipleSum => Box::new(MultipleSumUpdater::new()), + AggregationType::MultipleIncrease => Box::new(MultipleIncreaseUpdater::new()), + AggregationType::MultipleMinMax => Box::new(MultipleMinMaxUpdater::new( if sub_type.eq_ignore_ascii_case("max") { "max".to_string() } else { "min".to_string() }, )), - "Sum" | "sum" => Box::new(SumAccumulatorUpdater::new()), - "Min" | "min" => Box::new(MinMaxAccumulatorUpdater::new("min".to_string())), - "Max" | "max" => Box::new(MinMaxAccumulatorUpdater::new("max".to_string())), - "Increase" | "increase" => Box::new(IncreaseAccumulatorUpdater::new()), - "CountMinSketch" | "count_min_sketch" | "CMS" | "cms" => { + AggregationType::Sum => Box::new(SumAccumulatorUpdater::new()), + AggregationType::MinMax => Box::new(MinMaxAccumulatorUpdater::new( + if sub_type.eq_ignore_ascii_case("max") { + "max".to_string() + } else { + "min".to_string() + }, + )), + AggregationType::Increase => Box::new(IncreaseAccumulatorUpdater::new()), + AggregationType::CountMinSketch | AggregationType::CountMinSketchWithHeap => { let (row_num, col_num) = cms_params(config); Box::new(CmsAccumulatorUpdater::new(row_num, col_num)) } - "HydraKLL" | "hydra_kll" => { + AggregationType::HydraKLL => { let (row_num, col_num, k) = hydra_kll_params(config); Box::new(HydraKllAccumulatorUpdater::new(row_num, col_num, k)) } other => { tracing::warn!( - "Unknown aggregation_type '{}', defaulting to SingleSubpopulation Sum", + "Unknown aggregation_type '{:?}', defaulting to SingleSubpopulation Sum", other ); Box::new(SumAccumulatorUpdater::new()) @@ -649,7 +646,7 @@ pub fn create_accumulator_updater(config: &AggregationConfig) -> Box &'static str { - "CountMinSketchAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountMinSketch } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs index 1a2c827..f2751c3 100644 --- a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ - AggregateCore, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, - SerializableToSink, + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, }; use serde_json::Value; use sketch_core::count_min_with_heap::{CountMinSketchWithHeap, HeapItem}; @@ -171,8 +171,8 @@ impl AggregateCore for CountMinSketchWithHeapAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "CountMinSketchWithHeapAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountMinSketchWithHeap } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs index 7297680..1781de0 100644 --- a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs @@ -1,5 +1,6 @@ use crate::data_model::{ - AggregateCore, MergeableAccumulator, SerializableToSink, SingleSubpopulationAggregate, + AggregateCore, AggregationType, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, }; use base64::{engine::general_purpose, Engine as _}; use serde_json::Value; @@ -79,9 +80,9 @@ impl DatasketchesKLLAccumulator { let mut kll_accumulators = Vec::with_capacity(accumulators.len()); for acc in accumulators { - if acc.get_accumulator_type() != "DatasketchesKLLAccumulator" { + if acc.get_accumulator_type() != AggregationType::DatasketchesKLL { return Err(format!( - "Cannot merge DatasketchesKLLAccumulator with {}", + "Cannot merge DatasketchesKLLAccumulator with {:?}", acc.get_accumulator_type() ) .into()); @@ -192,8 +193,8 @@ impl AggregateCore for DatasketchesKLLAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "DatasketchesKLLAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::DatasketchesKLL } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs index d12c848..f0ae9f9 100644 --- a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ - AggregateCore, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, - SerializableToSink, + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, }; use serde_json::Value; use sketch_core::delta_set_aggregator::{deserialize_msgpack, serialize_msgpack}; @@ -240,8 +240,8 @@ impl AggregateCore for DeltaSetAggregatorAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "DeltaSetAggregatorAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::DeltaSetAggregator } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs index 42a9489..96cf19d 100644 --- a/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs @@ -1,6 +1,7 @@ use crate::{ data_model::{ - AggregateCore, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, + AggregateCore, AggregationType, MergeableAccumulator, MultipleSubpopulationAggregate, + SerializableToSink, }, KeyByLabelValues, }; @@ -108,8 +109,8 @@ impl AggregateCore for HydraKllSketchAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "HydraKllSketchAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::HydraKLL } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/increase_accumulator.rs b/asap-query-engine/src/precompute_operators/increase_accumulator.rs index 9dd5af0..5fcbe01 100644 --- a/asap-query-engine/src/precompute_operators/increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/increase_accumulator.rs @@ -1,5 +1,5 @@ use crate::data_model::{ - AggregateCore, Measurement, MergeableAccumulator, SerializableToSink, + AggregateCore, AggregationType, Measurement, MergeableAccumulator, SerializableToSink, SingleSubpopulationAggregate, SingleSubpopulationAggregateFactory, }; use serde::{Deserialize, Serialize}; @@ -241,8 +241,8 @@ impl AggregateCore for IncreaseAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "IncreaseAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Increase } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/min_max_accumulator.rs b/asap-query-engine/src/precompute_operators/min_max_accumulator.rs index 3b8361f..3b70087 100644 --- a/asap-query-engine/src/precompute_operators/min_max_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/min_max_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ - AggregateCore, MergeableAccumulator, SerializableToSink, SingleSubpopulationAggregate, - SingleSubpopulationAggregateFactory, + AggregateCore, AggregationType, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, SingleSubpopulationAggregateFactory, }; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -180,8 +180,8 @@ impl AggregateCore for MinMaxAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "MinMaxAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::MinMax } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs index f81fc60..82b77bb 100644 --- a/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ - AggregateCore, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, - SerializableToSink, SingleSubpopulationAggregate, + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, SingleSubpopulationAggregate, }; use crate::precompute_operators::IncreaseAccumulator; use serde::{Deserialize, Serialize}; @@ -278,8 +278,8 @@ impl AggregateCore for MultipleIncreaseAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "MultipleIncreaseAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::MultipleIncrease } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs index d8533e7..a8f14a3 100644 --- a/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ - AggregateCore, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, - SerializableToSink, + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, }; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -236,8 +236,8 @@ impl AggregateCore for MultipleMinMaxAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "MultipleMinMaxAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::MultipleMinMax } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs index c1d24af..c2ce83a 100644 --- a/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ - AggregateCore, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, - MultipleSubpopulationAggregateFactory, SerializableToSink, + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, MultipleSubpopulationAggregateFactory, SerializableToSink, }; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -226,8 +226,8 @@ impl AggregateCore for MultipleSumAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "MultipleSumAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::MultipleSum } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs b/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs index 6f38007..e93ae66 100644 --- a/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ - AggregateCore, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, - SerializableToSink, + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, }; use serde_json::Value; use sketch_core::set_aggregator::SetAggregator; @@ -175,8 +175,8 @@ impl AggregateCore for SetAggregatorAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "SetAggregatorAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::SetAggregator } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/sum_accumulator.rs b/asap-query-engine/src/precompute_operators/sum_accumulator.rs index 66f03a4..717bcd6 100644 --- a/asap-query-engine/src/precompute_operators/sum_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/sum_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ - AggregateCore, MergeableAccumulator, SerializableToSink, SingleSubpopulationAggregate, - SingleSubpopulationAggregateFactory, + AggregateCore, AggregationType, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, SingleSubpopulationAggregateFactory, }; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -114,8 +114,8 @@ impl AggregateCore for SumAccumulator { Ok(Box::new(merged)) } - fn get_accumulator_type(&self) -> &'static str { - "SumAccumulator" + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Sum } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/stores/simple_map_store/global.rs b/asap-query-engine/src/stores/simple_map_store/global.rs index 37ca9a6..66529db 100644 --- a/asap-query-engine/src/stores/simple_map_store/global.rs +++ b/asap-query-engine/src/stores/simple_map_store/global.rs @@ -1,4 +1,6 @@ -use crate::data_model::{AggregateCore, CleanupPolicy, PrecomputedOutput, StreamingConfig}; +use crate::data_model::{ + AggregateCore, AggregationType, CleanupPolicy, PrecomputedOutput, StreamingConfig, +}; use crate::stores::simple_map_store::common::{ EpochID, InternTable, MetricBucketMap, MutableEpoch, SealedEpoch, TimestampRange, }; @@ -218,7 +220,8 @@ impl Store for SimpleMapStoreGlobal { ( BatchConfig { metric: aggregation_config.metric.clone(), - is_delta: aggregation_config.aggregation_type == "DeltaSetAggregator", + is_delta: aggregation_config.aggregation_type + == AggregationType::DeltaSetAggregator, num_aggregates_to_retain: aggregation_config.num_aggregates_to_retain, read_count_threshold: aggregation_config.read_count_threshold, }, diff --git a/asap-query-engine/src/stores/simple_map_store/legacy/global.rs b/asap-query-engine/src/stores/simple_map_store/legacy/global.rs index 5d842e0..476c266 100644 --- a/asap-query-engine/src/stores/simple_map_store/legacy/global.rs +++ b/asap-query-engine/src/stores/simple_map_store/legacy/global.rs @@ -1,5 +1,6 @@ use crate::data_model::{ - AggregateCore, CleanupPolicy, KeyByLabelValues, PrecomputedOutput, StreamingConfig, + AggregateCore, AggregationType, CleanupPolicy, KeyByLabelValues, PrecomputedOutput, + StreamingConfig, }; use crate::stores::{Store, StoreResult, TimestampedBucketsMap}; use std::collections::HashMap; @@ -280,7 +281,7 @@ impl Store for LegacySimpleMapStoreGlobal { store_value.push((output.key, precompute)); // Apply retention policy if configured (but exclude DeltaSetAggregator) - if aggregation_config.aggregation_type != "DeltaSetAggregator" { + if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator { self.cleanup_old_aggregates( &mut data, &metric, diff --git a/asap-query-engine/src/stores/simple_map_store/legacy/per_key.rs b/asap-query-engine/src/stores/simple_map_store/legacy/per_key.rs index 8f6745c..b7807e1 100644 --- a/asap-query-engine/src/stores/simple_map_store/legacy/per_key.rs +++ b/asap-query-engine/src/stores/simple_map_store/legacy/per_key.rs @@ -1,5 +1,6 @@ use crate::data_model::{ - AggregateCore, CleanupPolicy, KeyByLabelValues, PrecomputedOutput, StreamingConfig, + AggregateCore, AggregationType, CleanupPolicy, KeyByLabelValues, PrecomputedOutput, + StreamingConfig, }; use crate::stores::{Store, StoreResult, TimestampedBucketsMap}; use dashmap::DashMap; @@ -276,7 +277,7 @@ impl LegacySimpleMapStorePerKey { .get_aggregation_config(aggregation_id) .ok_or_else(|| format!("Aggregation config not found for {}", aggregation_id))?; - if aggregation_config.aggregation_type != "DeltaSetAggregator" { + if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator { self.cleanup_old_aggregates( &mut data, metric, diff --git a/asap-query-engine/src/stores/simple_map_store/per_key.rs b/asap-query-engine/src/stores/simple_map_store/per_key.rs index 5a6cbd3..a1cdfb8 100644 --- a/asap-query-engine/src/stores/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/simple_map_store/per_key.rs @@ -1,4 +1,6 @@ -use crate::data_model::{AggregateCore, CleanupPolicy, PrecomputedOutput, StreamingConfig}; +use crate::data_model::{ + AggregateCore, AggregationType, CleanupPolicy, PrecomputedOutput, StreamingConfig, +}; use crate::stores::simple_map_store::common::{ EpochID, InternTable, MetricBucketMap, MetricID, MutableEpoch, SealedEpoch, TimestampRange, }; @@ -305,7 +307,7 @@ impl SimpleMapStorePerKey { .ok_or_else(|| format!("Aggregation config not found for {}", aggregation_id))?; // Configure epoch capacity on first insert (Optimization 2) - if aggregation_config.aggregation_type != "DeltaSetAggregator" { + if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator { data.configure_epochs(aggregation_config.num_aggregates_to_retain); } @@ -319,7 +321,7 @@ impl SimpleMapStorePerKey { .insert(metric_id, timestamp_range, Arc::from(precompute)); // After each item, check if we should rotate (CircularBuffer, Optimization 2) - if aggregation_config.aggregation_type != "DeltaSetAggregator" + if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator && matches!(self.cleanup_policy, CleanupPolicy::CircularBuffer) { data.maybe_rotate_epoch(); @@ -327,7 +329,7 @@ impl SimpleMapStorePerKey { } // Apply retention policy if configured (but exclude DeltaSetAggregator) - if aggregation_config.aggregation_type != "DeltaSetAggregator" { + if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator { self.cleanup_old_aggregates( &mut data, metric, diff --git a/asap-query-engine/src/tests/capability_matching_tests.rs b/asap-query-engine/src/tests/capability_matching_tests.rs index 9bec8ad..db10b4c 100644 --- a/asap-query-engine/src/tests/capability_matching_tests.rs +++ b/asap-query-engine/src/tests/capability_matching_tests.rs @@ -5,8 +5,9 @@ //! the existing query_config path still takes priority when an entry is present. use crate::data_model::{ - AggregationConfig, AggregationReference, CleanupPolicy, InferenceConfig, PrecomputedOutput, - PromQLSchema, QueryConfig, QueryLanguage, SchemaConfig, StreamingConfig, WindowType, + AggregationConfig, AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, + PrecomputedOutput, PromQLSchema, QueryConfig, QueryLanguage, SchemaConfig, StreamingConfig, + WindowType, }; use crate::engines::simple_engine::SimpleEngine; use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; @@ -25,14 +26,14 @@ use std::sync::Arc; fn make_agg_config( id: u64, metric: &str, - agg_type: &str, + agg_type: AggregationType, window_size_s: u64, window_type: WindowType, grouping: &[&str], ) -> AggregationConfig { AggregationConfig { aggregation_id: id, - aggregation_type: agg_type.to_string(), + aggregation_type: agg_type, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(grouping.iter().map(|s| s.to_string()).collect()), @@ -161,7 +162,14 @@ fn engine_with_query_config( /// capability matching should route to it and return a valid context. #[test] fn capability_fallback_fires_when_no_config() { - let agg = make_agg_config(1, "cpu", "Sum", 300, WindowType::Tumbling, &[]); + let agg = make_agg_config( + 1, + "cpu", + AggregationType::Sum, + 300, + WindowType::Tumbling, + &[], + ); let engine = engine_no_query_configs("cpu", &[], vec![agg]); // sum_over_time(cpu[5m]) — 5 min = 300 s matches the 300 s tumbling config @@ -178,7 +186,14 @@ fn capability_fallback_fires_when_no_config() { /// We verify by giving the config a different agg_id than any compatible-by-type config. #[test] fn config_path_takes_priority_over_capability_matching() { - let agg = make_agg_config(42, "cpu", "Sum", 300, WindowType::Tumbling, &[]); + let agg = make_agg_config( + 42, + "cpu", + AggregationType::Sum, + 300, + WindowType::Tumbling, + &[], + ); let engine = engine_with_query_config("cpu", &[], agg, "sum_over_time(cpu[5m])"); let ctx = engine @@ -196,7 +211,7 @@ fn quantile_different_values_resolve_to_same_aggregation() { let kll = make_agg_config( 7, "latency", - "DatasketchesKLL", + AggregationType::DatasketchesKLL, 300, WindowType::Tumbling, &[], @@ -231,7 +246,14 @@ fn quantile_different_values_resolve_to_same_aggregation() { #[test] fn no_match_returns_none() { // KLL config present, but query asks for Sum — incompatible - let kll = make_agg_config(1, "cpu", "DatasketchesKLL", 300, WindowType::Tumbling, &[]); + let kll = make_agg_config( + 1, + "cpu", + AggregationType::DatasketchesKLL, + 300, + WindowType::Tumbling, + &[], + ); let engine = engine_no_query_configs("cpu", &[], vec![kll]); let ctx = @@ -245,8 +267,22 @@ fn no_match_returns_none() { /// When multiple compatible aggregations exist, the largest window should be preferred. #[test] fn priority_largest_window_wins() { - let small = make_agg_config(1, "cpu", "Sum", 300, WindowType::Tumbling, &[]); - let large = make_agg_config(2, "cpu", "Sum", 900, WindowType::Tumbling, &[]); + let small = make_agg_config( + 1, + "cpu", + AggregationType::Sum, + 300, + WindowType::Tumbling, + &[], + ); + let large = make_agg_config( + 2, + "cpu", + AggregationType::Sum, + 900, + WindowType::Tumbling, + &[], + ); let engine = engine_no_query_configs("cpu", &[], vec![small, large]); // sum_over_time(cpu[15m]) = 900 s — both 300 s and 900 s configs match (900 = 3×300), diff --git a/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs b/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs index 401f09f..7565717 100644 --- a/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs +++ b/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs @@ -17,7 +17,7 @@ mod tests { SetAggregatorAccumulator, SumAccumulator, }; use datafusion_summary_library::SketchType; - use promql_utilities::query_logics::enums::Statistic; + use promql_utilities::query_logics::enums::{AggregationType, Statistic}; use std::collections::HashMap; // ======================================================================== @@ -29,7 +29,7 @@ mod tests { let acc = SumAccumulator::with_sum(42.5); let bytes = serialize_accumulator_arroyo(&acc); let restored = deserialize_accumulator(&bytes, &SketchType::Sum).unwrap(); - assert_eq!(restored.get_accumulator_type(), "SumAccumulator"); + assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); // Query the restored accumulator via single subpopulation let restored_single = deserialize_single_subpopulation(&bytes, &SketchType::Sum).unwrap(); @@ -48,7 +48,7 @@ mod tests { let restored = deserialize_accumulator(&bytes, &SketchType::KLL).unwrap(); assert_eq!( restored.get_accumulator_type(), - "DatasketchesKLLAccumulator" + AggregationType::DatasketchesKLL ); // Query quantile via single subpopulation @@ -154,7 +154,7 @@ mod tests { let bytes = serialize_accumulator_arroyo(&hydra); let restored = deserialize_accumulator(&bytes, &SketchType::HydraKLL).unwrap(); - assert_eq!(restored.get_accumulator_type(), "HydraKllSketchAccumulator"); + assert_eq!(restored.get_accumulator_type(), AggregationType::HydraKLL); } #[test] @@ -167,13 +167,16 @@ mod tests { let bytes = serialize_accumulator_arroyo(&cms); let restored = deserialize_accumulator(&bytes, &SketchType::CountMinSketch).unwrap(); - assert_eq!(restored.get_accumulator_type(), "CountMinSketchAccumulator"); + assert_eq!( + restored.get_accumulator_type(), + AggregationType::CountMinSketch + ); let restored = deserialize_multiple_subpopulation(&bytes, &SketchType::CountMinSketch).unwrap(); assert_eq!( restored.clone_boxed().as_ref().get_accumulator_type(), - "CountMinSketchAccumulator" + AggregationType::CountMinSketch ); } @@ -306,7 +309,7 @@ mod tests { // Verify the arroyo bytes can be deserialized let restored = deserialize_accumulator(&arroyo_bytes, &SketchType::Sum).unwrap(); - assert_eq!(restored.get_accumulator_type(), "SumAccumulator"); + assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); } #[test] diff --git a/asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs b/asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs index c020c7d..69c888e 100644 --- a/asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs +++ b/asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs @@ -6,6 +6,7 @@ #[cfg(test)] mod tests { + use crate::data_model::AggregationType; use crate::precompute_operators::sum_accumulator::SumAccumulator; use crate::tests::test_utilities::engine_factories::{ create_engine_single_pop, create_engine_two_metrics, @@ -17,7 +18,7 @@ mod tests { async fn test_handle_query_promql_binary_returns_result() { let engine = create_engine_two_metrics( "errors_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![( Some(vec!["host-a".to_string()]), @@ -25,7 +26,7 @@ mod tests { )], "sum(errors_total) by (host)", "requests_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![( Some(vec!["host-a".to_string()]), @@ -54,7 +55,7 @@ mod tests { // Only requests_total is configured; foo() is not a known function. let engine = create_engine_single_pop( "requests_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![( Some(vec!["host-a".to_string()]), @@ -78,7 +79,7 @@ mod tests { async fn test_handle_query_promql_scalar_binary_returns_result() { let engine = create_engine_single_pop( "errors_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![( Some(vec!["host-a".to_string()]), @@ -104,7 +105,7 @@ mod tests { // Regression: single-metric queries continue to work after binary dispatch is wired in. let engine = create_engine_single_pop( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![ ( diff --git a/asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs b/asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs index cdcddab..c663906 100644 --- a/asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs @@ -13,7 +13,7 @@ mod tests { use datafusion::logical_expr::LogicalPlan; use promql_parser::parser::token::{TokenType, T_ADD, T_DIV, T_MOD, T_MUL, T_POW, T_SUB}; use promql_utilities::data_model::KeyByLabelNames; - use promql_utilities::query_logics::enums::Statistic; + use promql_utilities::query_logics::enums::{AggregationType, Statistic}; use std::collections::HashMap; fn make_context( @@ -42,8 +42,8 @@ mod tests { agg_info: AggregationIdInfo { aggregation_id_for_key: 1, aggregation_id_for_value: 1, - aggregation_type_for_key: "SumAggregator".to_string(), - aggregation_type_for_value: "SumAggregator".to_string(), + aggregation_type_for_key: AggregationType::Sum, + aggregation_type_for_value: AggregationType::Sum, }, do_merge: false, spatial_filter: String::new(), diff --git a/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs b/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs index f16835b..8764a80 100644 --- a/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs @@ -10,12 +10,12 @@ mod tests { QueryExecutionContext, QueryMetadata, StoreQueryParams, StoreQueryPlan, }; use promql_utilities::data_model::KeyByLabelNames; - use promql_utilities::query_logics::enums::Statistic; + use promql_utilities::query_logics::enums::{AggregationType, Statistic}; use std::collections::HashMap; fn create_context( statistic: Statistic, - aggregation_type: &str, + aggregation_type: AggregationType, output_labels: Vec<&str>, kwargs: HashMap, ) -> QueryExecutionContext { @@ -43,8 +43,8 @@ mod tests { agg_info: AggregationIdInfo { aggregation_id_for_key: 1, aggregation_id_for_value: 1, - aggregation_type_for_key: String::new(), - aggregation_type_for_value: aggregation_type.to_string(), + aggregation_type_for_key: AggregationType::Sum, + aggregation_type_for_value: aggregation_type, }, do_merge: false, spatial_filter: String::new(), @@ -61,15 +61,15 @@ mod tests { #[test] fn test_all_statistics_map_without_panic() { let statistics = vec![ - (Statistic::Sum, "SumAccumulator"), - (Statistic::Min, "MinMaxAccumulator"), - (Statistic::Max, "MinMaxAccumulator"), - (Statistic::Count, "SumAccumulator"), - (Statistic::Increase, "IncreaseAccumulator"), - (Statistic::Rate, "IncreaseAccumulator"), - (Statistic::Quantile, "DatasketchesKLLAccumulator"), - (Statistic::Cardinality, "SetAggregator"), - (Statistic::Topk, "CountMinSketchAccumulator"), + (Statistic::Sum, AggregationType::Sum), + (Statistic::Min, AggregationType::MinMax), + (Statistic::Max, AggregationType::MinMax), + (Statistic::Count, AggregationType::Sum), + (Statistic::Increase, AggregationType::Increase), + (Statistic::Rate, AggregationType::Increase), + (Statistic::Quantile, AggregationType::DatasketchesKLL), + (Statistic::Cardinality, AggregationType::SetAggregator), + (Statistic::Topk, AggregationType::CountMinSketch), ]; for (stat, agg_type) in statistics { @@ -96,7 +96,7 @@ mod tests { let ctx = create_context( Statistic::Topk, - "CountMinSketchAccumulator", + AggregationType::CountMinSketch, vec!["host"], kwargs, ); @@ -111,7 +111,7 @@ mod tests { use datafusion_summary_library::InferOperation; let ctx = create_context( Statistic::Topk, - "CountMinSketchAccumulator", + AggregationType::CountMinSketch, vec!["host"], HashMap::new(), ); @@ -130,7 +130,7 @@ mod tests { use datafusion_summary_library::InferOperation; let ctx = create_context( Statistic::Cardinality, - "SetAggregator", + AggregationType::SetAggregator, vec!["host"], HashMap::new(), ); @@ -145,7 +145,7 @@ mod tests { use datafusion_summary_library::InferOperation; let ctx = create_context( Statistic::Rate, - "IncreaseAccumulator", + AggregationType::Increase, vec!["host"], HashMap::new(), ); @@ -160,7 +160,7 @@ mod tests { use datafusion_summary_library::InferOperation; let ctx = create_context( Statistic::Count, - "SumAccumulator", + AggregationType::Sum, vec!["host"], HashMap::new(), ); @@ -176,18 +176,19 @@ mod tests { #[test] fn test_unknown_agg_type_errors() { + // SingleSubpopulation is a legacy wrapper variant not mapped to a SketchType let ctx = create_context( Statistic::Sum, - "FooBarAccumulator", + AggregationType::SingleSubpopulation, vec!["host"], HashMap::new(), ); let result = ctx.to_logical_plan(); - assert!(result.is_err(), "Unknown aggregation type should error"); + assert!(result.is_err(), "Unmapped aggregation type should error"); let err_msg = format!("{}", result.unwrap_err()); assert!( - err_msg.contains("FooBarAccumulator") || err_msg.contains("Unknown"), - "Error should mention the unknown type, got: {}", + err_msg.contains("Unknown"), + "Error should mention Unknown, got: {}", err_msg ); } @@ -198,7 +199,7 @@ mod tests { #[test] fn test_plan_with_empty_labels() { - let ctx = create_context(Statistic::Sum, "SumAccumulator", vec![], HashMap::new()); + let ctx = create_context(Statistic::Sum, AggregationType::Sum, vec![], HashMap::new()); let result = ctx.to_logical_plan(); assert!( result.is_ok(), diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs index b196841..91d1356 100644 --- a/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs @@ -6,6 +6,7 @@ #[cfg(test)] mod tests { + use crate::data_model::AggregationType; use crate::precompute_operators::sum_accumulator::SumAccumulator; use crate::tests::test_utilities::engine_factories::{ create_engine_three_metrics, create_engine_two_metrics, @@ -49,12 +50,12 @@ mod tests { let (data_errors, data_requests) = host_a_b_data(100.0, 200.0); let engine = create_engine_two_metrics( "errors_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_errors, "sum(errors_total) by (host)", "requests_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_requests, "sum(requests_total) by (host)", @@ -85,12 +86,12 @@ mod tests { let (data_a, data_b) = host_a_b_data(3.0, 4.0); let engine = create_engine_two_metrics( "metric_a", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_a, "sum(metric_a) by (host)", "metric_b", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_b, "sum(metric_b) by (host)", @@ -116,12 +117,12 @@ mod tests { let (data_a, data_b) = host_a_b_data(10.0, 20.0); let engine = create_engine_two_metrics( "metric_a", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_a, "sum(metric_a) by (host)", "metric_b", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_b, "sum(metric_b) by (host)", @@ -147,12 +148,12 @@ mod tests { let (data_a, data_b) = host_a_b_data(50.0, 30.0); let engine = create_engine_two_metrics( "metric_a", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_a, "sum(metric_a) by (host)", "metric_b", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_b, "sum(metric_b) by (host)", @@ -193,12 +194,12 @@ mod tests { let engine = create_engine_two_metrics( "errors_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_errors, "sum(errors_total) by (host)", "requests_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_requests, "sum(requests_total) by (host)", @@ -229,13 +230,13 @@ mod tests { )]; let engine = create_engine_two_metrics( "errors_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, "sum(errors_total) by (host)", // second metric not used but factory requires it; use empty data "dummy", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![], "sum(dummy) by (host)", @@ -264,12 +265,12 @@ mod tests { )]; let engine = create_engine_two_metrics( "success_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, "sum(success_total) by (host)", "dummy", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![], "sum(dummy) by (host)", @@ -328,17 +329,17 @@ mod tests { let engine = create_engine_three_metrics( "metric_a", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_a, "sum(metric_a) by (host)", "metric_b", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_b, "sum(metric_b) by (host)", "metric_c", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data_c, "sum(metric_c) by (host)", diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs index c20ac2a..649b2b2 100644 --- a/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs @@ -11,7 +11,7 @@ #[cfg(test)] mod tests { - use crate::data_model::{KeyByLabelValues, Measurement}; + use crate::data_model::{AggregationType, KeyByLabelValues, Measurement}; use crate::precompute_operators::{ CountMinSketchAccumulator, DeltaSetAggregatorAccumulator, HydraKllSketchAccumulator, IncreaseAccumulator, MultipleIncreaseAccumulator, @@ -59,7 +59,7 @@ mod tests { let engine = create_engine_single_pop_with_aggregated( "http_requests_total", - "MultipleIncreaseAccumulator", + AggregationType::MultipleIncrease, vec![], vec!["host", "endpoint"], vec![(None, Box::new(acc))], @@ -88,7 +88,7 @@ mod tests { let engine = create_engine_single_pop_with_aggregated( "http_requests_total", - "MultipleIncreaseAccumulator", + AggregationType::MultipleIncrease, vec![], vec!["host", "service"], vec![(None, Box::new(acc))], @@ -145,8 +145,8 @@ mod tests { let engine = create_engine_dual_input( "request_duration", - "HydraKllSketchAccumulator", - "DeltaSetAggregator", + AggregationType::HydraKLL, + AggregationType::DeltaSetAggregator, vec![], // grouping_labels: no store GROUP BY vec!["host", "endpoint"], // aggregated_labels: sub-keys tracked by DeltaSet vec![(None, Box::new(hydra))], // store key = None @@ -187,8 +187,8 @@ mod tests { let engine = create_engine_dual_input( "event_frequency", - "CountMinSketchAccumulator", - "DeltaSetAggregator", + AggregationType::CountMinSketch, + AggregationType::DeltaSetAggregator, vec![], // grouping_labels: no store GROUP BY vec!["host", "event"], // aggregated_labels: sub-keys tracked by DeltaSet vec![(None, Box::new(cms))], // store key = None @@ -220,7 +220,7 @@ mod tests { let engine = create_engine_single_pop_with_aggregated( "requests", - "MultipleIncreaseAccumulator", + AggregationType::MultipleIncrease, vec![], vec!["host", "endpoint"], vec![(None, Box::new(acc1)), (None, Box::new(acc2))], @@ -249,7 +249,7 @@ mod tests { let engine = create_engine_single_pop_with_aggregated( "requests", - "MultipleIncreaseAccumulator", + AggregationType::MultipleIncrease, vec![], vec!["host", "endpoint"], vec![(None, Box::new(empty))], @@ -289,8 +289,8 @@ mod tests { let engine = create_engine_dual_input( "request_duration", - "HydraKllSketchAccumulator", - "DeltaSetAggregator", + AggregationType::HydraKLL, + AggregationType::DeltaSetAggregator, vec![], // no store GROUP BY vec!["host", "endpoint"], // aggregated_labels vec![(None, Box::new(hydra))], diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs index 5b0d5e1..8d7e915 100644 --- a/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs @@ -10,7 +10,7 @@ use std::collections::HashMap; #[cfg(test)] mod tests { use super::*; - use crate::data_model::WindowType; + use crate::data_model::{AggregationType, WindowType}; use crate::precompute_operators::DatasketchesKLLAccumulator; use crate::tests::test_utilities::engine_factories::*; @@ -47,7 +47,7 @@ mod tests { let query = "sum_over_time(http_requests[5s])"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, query, @@ -77,7 +77,7 @@ mod tests { let query = "sum_over_time(http_requests[5s])"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, query, @@ -114,7 +114,7 @@ mod tests { let query = "sum_over_time(http_requests[5s])"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, query, @@ -151,7 +151,7 @@ mod tests { let query = "quantile_over_time(0.5, latency[5s])"; let engine = create_engine_multi_timestamp_with_window( "latency", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, vec!["host"], data, query, @@ -189,7 +189,7 @@ mod tests { let query = "sum_over_time(http_requests[5s])"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, query, @@ -222,7 +222,7 @@ mod tests { let query = "sum_over_time(http_requests[5s])"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![], query, @@ -243,7 +243,7 @@ mod tests { let query = "sum_over_time(http_requests[5s])"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![( 1_000_000, @@ -267,7 +267,7 @@ mod tests { let query = "sum_over_time(http_requests[5s])"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![( 1_000_000, @@ -320,7 +320,7 @@ mod tests { let query = "sum by (host) (sum_over_time(http_requests[5s]))"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, query, @@ -372,7 +372,7 @@ mod tests { let query = "sum by (host) (sum_over_time(http_requests[5s]))"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, query, @@ -407,7 +407,7 @@ mod tests { let query = "sum by (host) (sum_over_time(http_requests[5s]))"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![( 1_000_000, @@ -435,7 +435,7 @@ mod tests { let query = "sum by (host) (sum_over_time(http_requests[5s]))"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![( 1_000_000, @@ -463,7 +463,7 @@ mod tests { let query = "sum by (host) (sum_over_time(http_requests[5s]))"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![], query, @@ -499,7 +499,7 @@ mod tests { let query = "sum_over_time(http_requests[5s])"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, query, @@ -529,7 +529,7 @@ mod tests { let query = "sum by (host) (sum_over_time(http_requests[5s]))"; let engine = create_engine_multi_timestamp_with_window( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, query, diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs index 3431d06..ba32446 100644 --- a/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs @@ -5,7 +5,7 @@ //! //! These tests use an actual store with test data. -use crate::data_model::{KeyByLabelValues, Measurement}; +use crate::data_model::{AggregationType, KeyByLabelValues, Measurement}; use crate::engines::simple_engine::SimpleEngine; use crate::precompute_operators::sum_accumulator::SumAccumulator; use std::collections::HashMap; @@ -23,7 +23,7 @@ mod tests { fn create_test_engine_with_data() -> SimpleEngine { create_engine_single_pop( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![ ( @@ -132,7 +132,7 @@ mod tests { async fn test_execute_plan_multiple_timestamps() { let engine = create_engine_multi_timestamp( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![ ( @@ -192,7 +192,7 @@ mod tests { let engine = create_engine_single_pop( "latency", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, vec!["host"], vec![ (Some(vec!["host-a".to_string()]), Box::new(kll_a)), @@ -213,7 +213,7 @@ mod tests { let engine = create_engine_single_pop( "latency", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, vec!["host"], vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], "quantile(0.99, latency) by (host)", @@ -231,7 +231,7 @@ mod tests { let engine = create_engine_single_pop( "latency", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, vec!["host"], vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], "quantile(0.0, latency) by (host)", @@ -249,7 +249,7 @@ mod tests { let engine = create_engine_single_pop( "latency", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, vec!["host"], vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], "quantile(1.0, latency) by (host)", @@ -267,7 +267,7 @@ mod tests { let engine = create_engine_single_pop( "latency", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, vec!["host"], vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], "quantile(0.25, latency) by (host)", @@ -293,7 +293,7 @@ mod tests { let engine = create_engine_single_pop( "active_users", - "SetAggregator", + AggregationType::SetAggregator, vec!["host"], vec![ (Some(vec!["host-a".to_string()]), Box::new(set_a)), @@ -313,7 +313,7 @@ mod tests { let engine = create_engine_single_pop( "http_requests_total", - "IncreaseAccumulator", + AggregationType::Increase, vec!["host"], vec![ (Some(vec!["host-a".to_string()]), Box::new(inc_a)), @@ -338,7 +338,7 @@ mod tests { let engine = create_engine_single_pop( "temperature", - "MinMaxAccumulator", + AggregationType::MinMax, vec!["host"], vec![ (Some(vec!["host-a".to_string()]), Box::new(mm_a)), @@ -358,7 +358,7 @@ mod tests { let engine = create_engine_single_pop( "temperature", - "MinMaxAccumulator", + AggregationType::MinMax, vec!["host"], vec![ (Some(vec!["host-a".to_string()]), Box::new(mm_a)), @@ -380,7 +380,7 @@ mod tests { let engine = create_engine_single_pop( "requests", - "MultipleSumAccumulator", + AggregationType::MultipleSum, vec!["host"], vec![(Some(vec!["host-a".to_string()]), Box::new(ms))], "sum(requests) by (host)", @@ -396,7 +396,7 @@ mod tests { let engine = create_engine_single_pop( "latency", - "MultipleMinMaxAccumulator", + AggregationType::MultipleMinMax, vec!["host"], vec![(Some(vec!["host-a".to_string()]), Box::new(mm))], "min(latency) by (host)", @@ -413,7 +413,7 @@ mod tests { async fn test_execute_plan_empty_store() { let engine = create_engine_single_pop( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![], // No data "sum(http_requests) by (host)", @@ -441,7 +441,7 @@ mod tests { let engine = create_engine_single_pop( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, "sum(http_requests) by (host)", @@ -472,7 +472,7 @@ mod tests { let engine = create_engine_multi_timestamp( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, "sum(http_requests) by (host)", @@ -516,7 +516,7 @@ mod tests { let engine = create_engine_multi_timestamp( "http_requests", - "SumAccumulator", + AggregationType::Sum, vec!["host"], data, "sum(http_requests) by (host)", @@ -542,7 +542,7 @@ mod tests { let engine = create_engine_single_pop( "http_requests_total", - "IncreaseAccumulator", + AggregationType::Increase, vec!["host"], vec![(Some(vec!["host-a".to_string()]), Box::new(inc))], "sum(increase(http_requests_total[10s])) by (host)", @@ -568,7 +568,7 @@ mod tests { let engine = create_engine_single_pop( "temperature", - "MinMaxAccumulator", + AggregationType::MinMax, vec!["host"], vec![(Some(vec!["host-a".to_string()]), Box::new(mm))], "min(temperature) by (host)", diff --git a/asap-query-engine/src/tests/datafusion/structural_matching_tests.rs b/asap-query-engine/src/tests/datafusion/structural_matching_tests.rs index 37fecb8..d04d99a 100644 --- a/asap-query-engine/src/tests/datafusion/structural_matching_tests.rs +++ b/asap-query-engine/src/tests/datafusion/structural_matching_tests.rs @@ -6,6 +6,7 @@ #[cfg(test)] mod tests { + use crate::data_model::AggregationType; use crate::precompute_operators::sum_accumulator::SumAccumulator; use crate::tests::test_utilities::engine_factories::create_engine_single_pop; @@ -13,7 +14,7 @@ mod tests { fn test_structural_match_rate_query_finds_config() { let engine = create_engine_single_pop( "http_requests_total", - "MultipleIncrease", + AggregationType::MultipleIncrease, vec!["host"], vec![( Some(vec!["host-a".to_string()]), @@ -35,7 +36,7 @@ mod tests { fn test_structural_match_wrong_metric_returns_none() { let engine = create_engine_single_pop( "http_requests_total", - "MultipleIncrease", + AggregationType::MultipleIncrease, vec!["host"], vec![( Some(vec!["host-a".to_string()]), @@ -56,7 +57,7 @@ mod tests { fn test_structural_match_wrong_range_returns_none() { let engine = create_engine_single_pop( "http_requests_total", - "MultipleIncrease", + AggregationType::MultipleIncrease, vec!["host"], vec![( Some(vec!["host-a".to_string()]), @@ -78,7 +79,7 @@ mod tests { fn test_structural_match_wrong_function_returns_none() { let engine = create_engine_single_pop( "http_requests_total", - "MultipleIncrease", + AggregationType::MultipleIncrease, vec!["host"], vec![( Some(vec!["host-a".to_string()]), @@ -100,7 +101,7 @@ mod tests { fn test_structural_match_spatial_query() { let engine = create_engine_single_pop( "http_requests_total", - "SumAccumulator", + AggregationType::Sum, vec!["host"], vec![( Some(vec!["host-a".to_string()]), diff --git a/asap-query-engine/src/tests/elastic_dsl_query_tests.rs b/asap-query-engine/src/tests/elastic_dsl_query_tests.rs index 09c7ac5..07097ea 100644 --- a/asap-query-engine/src/tests/elastic_dsl_query_tests.rs +++ b/asap-query-engine/src/tests/elastic_dsl_query_tests.rs @@ -5,7 +5,7 @@ #[cfg(test)] mod tests { - use crate::data_model::{AggregateCore, KeyByLabelValues}; + use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues}; use crate::precompute_operators::DatasketchesKLLAccumulator; use crate::tests::test_utilities::create_engine_multi_timestamp; use crate::QueryResult; @@ -72,7 +72,7 @@ mod tests { let engine = create_engine_multi_timestamp( "http_requests", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, Vec::new(), // No labels for this test kll_data, &elastic_query.to_string(), @@ -135,7 +135,7 @@ mod tests { let engine = create_engine_multi_timestamp( "http_requests", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, vec!["host"], kll_data, &elastic_query.to_string(), @@ -217,7 +217,7 @@ mod tests { let engine = create_engine_multi_timestamp( "http_requests", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, vec!["host", "region"], kll_data, &elastic_query.to_string(), @@ -303,7 +303,7 @@ mod tests { let engine = create_engine_multi_timestamp( "http_requests", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, vec!["host", "region"], kll_data, &elastic_query.to_string(), @@ -355,7 +355,7 @@ mod tests { let engine = create_engine_multi_timestamp( "http_requests", - "DatasketchesKLLAccumulator", + AggregationType::DatasketchesKLL, Vec::new(), // No labels for this test kll_data, &elastic_query.to_string(), diff --git a/asap-query-engine/src/tests/sql_pattern_matching_tests.rs b/asap-query-engine/src/tests/sql_pattern_matching_tests.rs index e0e566c..44db4fc 100644 --- a/asap-query-engine/src/tests/sql_pattern_matching_tests.rs +++ b/asap-query-engine/src/tests/sql_pattern_matching_tests.rs @@ -6,8 +6,8 @@ #[cfg(test)] mod tests { use crate::data_model::{ - AggregationConfig, AggregationReference, CleanupPolicy, InferenceConfig, QueryConfig, - QueryLanguage, SchemaConfig, StreamingConfig, WindowType, + AggregationConfig, AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, + QueryConfig, QueryLanguage, SchemaConfig, StreamingConfig, WindowType, }; use crate::engines::simple_engine::SimpleEngine; use crate::stores::simple_map_store::SimpleMapStore; @@ -49,7 +49,7 @@ mod tests { // Streaming config let agg_config = AggregationConfig { aggregation_id: agg_id, - aggregation_type: "SumAccumulator".to_string(), + aggregation_type: AggregationType::Sum, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new( diff --git a/asap-query-engine/src/tests/store_correctness_tests.rs b/asap-query-engine/src/tests/store_correctness_tests.rs index 62f7215..2032aa0 100644 --- a/asap-query-engine/src/tests/store_correctness_tests.rs +++ b/asap-query-engine/src/tests/store_correctness_tests.rs @@ -28,8 +28,8 @@ //! | `contract_global` | `LockStrategy::Global` | use crate::data_model::{ - CleanupPolicy, KeyByLabelValues, LockStrategy, Measurement, SerializableToSink, - StreamingConfig, WindowType, + AggregationType, CleanupPolicy, KeyByLabelValues, LockStrategy, Measurement, + SerializableToSink, StreamingConfig, WindowType, }; use crate::precompute_operators::{ CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, DatasketchesKLLAccumulator, @@ -47,13 +47,13 @@ use std::sync::Arc; fn make_agg_config( agg_id: u64, - aggregation_type: &str, + aggregation_type: AggregationType, num_aggregates_to_retain: Option, read_count_threshold: Option, ) -> AggregationConfig { AggregationConfig::new( agg_id, - aggregation_type.to_string(), + aggregation_type, "".to_string(), HashMap::new(), KeyByLabelNames::empty(), @@ -72,7 +72,9 @@ fn make_agg_config( ) } -fn make_streaming_config(ids: &[(u64, &str, Option, Option)]) -> Arc { +fn make_streaming_config( + ids: &[(u64, AggregationType, Option, Option)], +) -> Arc { let configs = ids .iter() .map(|&(id, agg_type, retain, threshold)| { @@ -85,18 +87,18 @@ fn make_streaming_config(ids: &[(u64, &str, Option, Option)]) -> Arc, Option)], + ids: &[(u64, AggregationType, Option, Option)], ) -> SimpleMapStore { let config = make_streaming_config(ids); SimpleMapStore::new_with_strategy(config, policy, strategy) } -/// Convenience: single agg_id=1, type "Sum", no cleanup. +/// Convenience: single agg_id=1, type Sum, no cleanup. fn make_store_simple(strategy: LockStrategy) -> SimpleMapStore { make_store( strategy, CleanupPolicy::NoCleanup, - &[(1, "Sum", None, None)], + &[(1, AggregationType::Sum, None, None)], ) } @@ -414,7 +416,10 @@ fn test_multiple_agg_ids_are_isolated(strategy: LockStrategy) { let store = make_store( strategy, CleanupPolicy::NoCleanup, - &[(1, "Sum", None, None), (2, "Sum", None, None)], + &[ + (1, AggregationType::Sum, None, None), + (2, AggregationType::Sum, None, None), + ], ); let (o1, a1) = sum_entry(1, 1_000, 2_000, 10.0); let (o2, a2) = sum_entry(2, 3_000, 4_000, 20.0); @@ -475,7 +480,10 @@ fn test_earliest_timestamp_tracked_per_agg_id(strategy: LockStrategy) { let store = make_store( strategy, CleanupPolicy::NoCleanup, - &[(1, "Sum", None, None), (2, "Sum", None, None)], + &[ + (1, AggregationType::Sum, None, None), + (2, AggregationType::Sum, None, None), + ], ); let (o1, a1) = sum_entry(1, 1_000, 2_000, 1.0); let (o2, a2) = sum_entry(2, 9_000, 10_000, 1.0); @@ -505,7 +513,7 @@ fn test_cleanup_circular_buffer_evicts_oldest_window(strategy: LockStrategy) { let store = make_store( strategy, CleanupPolicy::CircularBuffer, - &[(1, "Sum", Some(2), None)], + &[(1, AggregationType::Sum, Some(2), None)], ); for i in 0u64..9 { let (out, acc) = sum_entry(1, i * 60_000, (i + 1) * 60_000, i as f64); @@ -525,7 +533,7 @@ fn test_cleanup_circular_buffer_retains_newest_windows(strategy: LockStrategy) { let store = make_store( strategy, CleanupPolicy::CircularBuffer, - &[(1, "Sum", Some(2), None)], + &[(1, AggregationType::Sum, Some(2), None)], ); for i in 0u64..9 { let (out, acc) = sum_entry(1, i * 60_000, (i + 1) * 60_000, i as f64); @@ -550,7 +558,7 @@ fn test_cleanup_read_based_evicts_after_threshold_reads(strategy: LockStrategy) let store = make_store( strategy, CleanupPolicy::ReadBased, - &[(1, "Sum", None, Some(2))], + &[(1, AggregationType::Sum, None, Some(2))], ); let (out, acc) = sum_entry(1, 1_000, 2_000, 1.0); store.insert_precomputed_output(out, acc).unwrap(); @@ -593,7 +601,7 @@ fn test_cleanup_read_based_unread_window_is_retained(strategy: LockStrategy) { let store = make_store( strategy, CleanupPolicy::ReadBased, - &[(1, "Sum", None, Some(1))], + &[(1, AggregationType::Sum, None, Some(1))], ); let (out, acc) = sum_entry(1, 1_000, 2_000, 1.0); store.insert_precomputed_output(out, acc).unwrap(); @@ -623,7 +631,7 @@ fn test_delta_set_aggregator_bypasses_cleanup(strategy: LockStrategy) { let store = make_store( strategy, CleanupPolicy::CircularBuffer, - &[(1, "DeltaSetAggregator", Some(2), None)], + &[(1, AggregationType::DeltaSetAggregator, Some(2), None)], ); let n = 10u64; for i in 0..n { diff --git a/asap-query-engine/src/tests/test_utilities/comparison.rs b/asap-query-engine/src/tests/test_utilities/comparison.rs index c5c5bdb..954d510 100644 --- a/asap-query-engine/src/tests/test_utilities/comparison.rs +++ b/asap-query-engine/src/tests/test_utilities/comparison.rs @@ -2,7 +2,7 @@ //! //! Provides assertion helpers for deep equality checking of query execution contexts. -use crate::data_model::AggregationIdInfo; +use crate::data_model::{AggregationIdInfo, AggregationType}; use crate::engines::simple_engine::{ QueryExecutionContext, QueryMetadata, StoreQueryParams, StoreQueryPlan, }; @@ -200,8 +200,8 @@ mod tests { agg_info: AggregationIdInfo { aggregation_id_for_key: 1, aggregation_id_for_value: 1, - aggregation_type_for_key: "SumAccumulator".to_string(), - aggregation_type_for_value: "SumAccumulator".to_string(), + aggregation_type_for_key: AggregationType::Sum, + aggregation_type_for_value: AggregationType::Sum, }, do_merge: true, // OnlyTemporal queries merge spatial_filter: String::new(), diff --git a/asap-query-engine/src/tests/test_utilities/config_builders.rs b/asap-query-engine/src/tests/test_utilities/config_builders.rs index 0183f8f..ec69127 100644 --- a/asap-query-engine/src/tests/test_utilities/config_builders.rs +++ b/asap-query-engine/src/tests/test_utilities/config_builders.rs @@ -4,8 +4,8 @@ //! and SQLSchema objects for testing. use crate::data_model::{ - AggregationConfig, AggregationReference, CleanupPolicy, InferenceConfig, PromQLSchema, - QueryConfig, SchemaConfig, StreamingConfig, WindowType, + AggregationConfig, AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, + PromQLSchema, QueryConfig, SchemaConfig, StreamingConfig, WindowType, }; use promql_utilities::data_model::KeyByLabelNames; use sql_utilities::sqlhelper::{SQLSchema, Table}; @@ -90,7 +90,7 @@ impl TestConfigBuilder { // Create streaming config for this aggregation let agg_config = AggregationConfig { aggregation_id: agg_id, - aggregation_type: "SumAccumulator".to_string(), + aggregation_type: AggregationType::Sum, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(self.grouping_labels.clone()), @@ -127,7 +127,7 @@ impl TestConfigBuilder { let agg_config = AggregationConfig { aggregation_id: agg_id, - aggregation_type: "SumAccumulator".to_string(), + aggregation_type: AggregationType::Sum, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(self.grouping_labels.clone()), @@ -170,7 +170,7 @@ impl TestConfigBuilder { let agg_config = AggregationConfig { aggregation_id: agg_id, - aggregation_type: "SumAccumulator".to_string(), // For collapsable sum queries + aggregation_type: AggregationType::Sum, // For collapsable sum queries aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(self.grouping_labels.clone()), diff --git a/asap-query-engine/src/tests/test_utilities/engine_factories.rs b/asap-query-engine/src/tests/test_utilities/engine_factories.rs index 51aaed2..2594721 100644 --- a/asap-query-engine/src/tests/test_utilities/engine_factories.rs +++ b/asap-query-engine/src/tests/test_utilities/engine_factories.rs @@ -6,9 +6,9 @@ //! the correct aggregation_type string. use crate::data_model::{ - AggregationConfig, AggregationReference, CleanupPolicy, InferenceConfig, KeyByLabelValues, - PrecomputedOutput, PromQLSchema, QueryConfig, QueryLanguage, SchemaConfig, StreamingConfig, - WindowType, + AggregationConfig, AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, + KeyByLabelValues, PrecomputedOutput, PromQLSchema, QueryConfig, QueryLanguage, SchemaConfig, + StreamingConfig, WindowType, }; use crate::engines::query_result::InstantVectorElement; use crate::engines::simple_engine::SimpleEngine; @@ -32,7 +32,7 @@ pub type AccumulatorData = Vec<(Option>, Box)>; /// * `promql_query` - The PromQL query string pub fn create_engine_single_pop( metric: &str, - aggregation_type: &str, + aggregation_type: AggregationType, grouping_labels: Vec<&str>, data: AccumulatorData, promql_query: &str, @@ -54,7 +54,7 @@ pub fn create_engine_single_pop( /// (e.g. "endpoint" within a MultipleIncrease accumulator). pub fn create_engine_single_pop_with_aggregated( metric: &str, - aggregation_type: &str, + aggregation_type: AggregationType, grouping_labels: Vec<&str>, aggregated_labels: Vec<&str>, data: AccumulatorData, @@ -73,7 +73,7 @@ pub fn create_engine_single_pop_with_aggregated( let mut aggregation_configs = HashMap::new(); let agg_config = AggregationConfig { aggregation_id: 1, - aggregation_type: aggregation_type.to_string(), + aggregation_type, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(grouping_label_strings.clone()), @@ -147,8 +147,8 @@ pub fn create_engine_single_pop_with_aggregated( #[allow(clippy::too_many_arguments)] pub fn create_engine_dual_input( metric: &str, - value_agg_type: &str, - key_agg_type: &str, + value_agg_type: AggregationType, + key_agg_type: AggregationType, grouping_labels: Vec<&str>, aggregated_labels: Vec<&str>, value_data: AccumulatorData, @@ -170,7 +170,7 @@ pub fn create_engine_dual_input( // Value aggregation (id=1) let value_agg_config = AggregationConfig { aggregation_id: 1, - aggregation_type: value_agg_type.to_string(), + aggregation_type: value_agg_type, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(grouping_label_strings.clone()), @@ -193,7 +193,7 @@ pub fn create_engine_dual_input( // Keys aggregation (id=2) let keys_agg_config = AggregationConfig { aggregation_id: 2, - aggregation_type: key_agg_type.to_string(), + aggregation_type: key_agg_type, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(grouping_label_strings.clone()), @@ -269,12 +269,12 @@ pub fn create_engine_dual_input( #[allow(clippy::too_many_arguments)] pub fn create_engine_two_metrics( metric_a: &str, - aggregation_type_a: &str, + aggregation_type_a: AggregationType, grouping_labels_a: Vec<&str>, data_a: AccumulatorData, query_a: &str, metric_b: &str, - aggregation_type_b: &str, + aggregation_type_b: AggregationType, grouping_labels_b: Vec<&str>, data_b: AccumulatorData, query_b: &str, @@ -286,7 +286,7 @@ pub fn create_engine_two_metrics( let agg_config_a = AggregationConfig { aggregation_id: 1, - aggregation_type: aggregation_type_a.to_string(), + aggregation_type: aggregation_type_a, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(labels_a.clone()), @@ -308,7 +308,7 @@ pub fn create_engine_two_metrics( let agg_config_b = AggregationConfig { aggregation_id: 2, - aggregation_type: aggregation_type_b.to_string(), + aggregation_type: aggregation_type_b, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(labels_b.clone()), @@ -381,17 +381,17 @@ pub fn create_engine_two_metrics( #[allow(clippy::too_many_arguments)] pub fn create_engine_three_metrics( metric_a: &str, - aggregation_type_a: &str, + aggregation_type_a: AggregationType, grouping_labels_a: Vec<&str>, data_a: AccumulatorData, query_a: &str, metric_b: &str, - aggregation_type_b: &str, + aggregation_type_b: AggregationType, grouping_labels_b: Vec<&str>, data_b: AccumulatorData, query_b: &str, metric_c: &str, - aggregation_type_c: &str, + aggregation_type_c: AggregationType, grouping_labels_c: Vec<&str>, data_c: AccumulatorData, query_c: &str, @@ -411,7 +411,7 @@ pub fn create_engine_three_metrics( id, AggregationConfig { aggregation_id: id, - aggregation_type: agg_type.to_string(), + aggregation_type: agg_type, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(labels.clone()), @@ -481,7 +481,7 @@ pub fn create_engine_three_metrics( #[allow(clippy::type_complexity)] pub fn create_engine_multi_timestamp( metric: &str, - aggregation_type: &str, + aggregation_type: AggregationType, grouping_labels: Vec<&str>, data: Vec<(u64, Option>, Box)>, promql_query: &str, @@ -492,7 +492,7 @@ pub fn create_engine_multi_timestamp( let mut aggregation_configs = HashMap::new(); let agg_config = AggregationConfig { aggregation_id: 1, - aggregation_type: aggregation_type.to_string(), + aggregation_type, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(grouping_label_strings.clone()), @@ -559,7 +559,7 @@ pub fn create_engine_multi_timestamp( #[allow(clippy::type_complexity)] pub fn create_engine_multi_timestamp_with_window( metric: &str, - aggregation_type: &str, + aggregation_type: AggregationType, grouping_labels: Vec<&str>, data: Vec<(u64, Option>, Box)>, promql_query: &str, @@ -572,7 +572,7 @@ pub fn create_engine_multi_timestamp_with_window( let mut aggregation_configs = HashMap::new(); let agg_config = AggregationConfig { aggregation_id: 1, - aggregation_type: aggregation_type.to_string(), + aggregation_type, aggregation_sub_type: String::new(), parameters: HashMap::new(), grouping_labels: KeyByLabelNames::new(grouping_label_strings.clone()), diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index 53dd572..c80f5ec 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -8,7 +8,7 @@ //! 4. Drains captured outputs and verifies equivalence with ArroYo-format accumulators use asap_types::aggregation_config::AggregationConfig; -use asap_types::enums::WindowType; +use asap_types::enums::{AggregationType, WindowType}; use flate2::{write::GzEncoder, Compression}; use prost::Message; use serde_json::json; @@ -32,7 +32,7 @@ use query_engine_rust::precompute_operators::multiple_sum_accumulator::MultipleS fn make_agg_config( id: u64, metric: &str, - agg_type: &str, + agg_type: AggregationType, agg_sub_type: &str, window_secs: u64, slide_secs: u64, @@ -45,7 +45,7 @@ fn make_agg_config( }; AggregationConfig::new( id, - agg_type.to_string(), + agg_type, agg_sub_type.to_string(), HashMap::new(), promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new( @@ -151,7 +151,7 @@ async fn e2e_kll_output_matches_arroyo() { let mut kll_config = make_agg_config( agg_id, "latency", - "DatasketchesKLL", + AggregationType::DatasketchesKLL, "", window_secs, 0, @@ -268,7 +268,7 @@ async fn e2e_multiple_sum_output_matches_arroyo() { let config = make_agg_config( agg_id, "cpu", - "MultipleSum", + AggregationType::MultipleSum, "sum", window_secs, 0,