From c2b63fca3c92d1314afae969a862b45a300b3d8c Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 20:44:14 -0400 Subject: [PATCH] refactor(query-engine): remove unused DataFusion query engine The logical/physical DataFusion plan engine (engines::logical, engines::physical) was superseded by a native implementation in #567. Its only entry points, execute_plan and execute_logical_plan on SimpleEngine, were both #[allow(dead_code)]: execute_logical_plan had zero callers anywhere, and execute_plan was reachable only from tests exercising the DataFusion path directly. Drop those modules, the demo bin that used them, and the DataFusion- specific tests, along with the datafusion/arrow/datafusion_summary_library dependencies they pulled in. Three test files under tests/datafusion/ that don't touch DataFusion (structural/range/dispatch arithmetic tests) move up to tests/ instead of being deleted. --- Cargo.lock | 3 - asap-query-engine/Cargo.toml | 3 - .../src/bin/show_logical_plans.rs | 528 --------- asap-query-engine/src/engines/logical/mod.rs | 6 - .../src/engines/logical/plan_builder.rs | 1014 ----------------- asap-query-engine/src/engines/mod.rs | 2 - .../src/engines/physical/accumulator_serde.rs | 363 ------ .../src/engines/physical/conversion.rs | 452 -------- asap-query-engine/src/engines/physical/mod.rs | 29 - .../src/engines/physical/planner.rs | 180 --- .../physical/precomputed_summary_read_exec.rs | 169 --- .../engines/physical/summary_infer_exec.rs | 769 ------------- .../physical/summary_merge_multiple_exec.rs | 556 --------- .../src/engines/simple_engine/mod.rs | 161 --- .../datafusion/accumulator_serde_tests.rs | 339 ------ asap-query-engine/src/tests/datafusion/mod.rs | 14 - .../datafusion/plan_builder_binary_tests.rs | 190 --- .../plan_builder_regression_tests.rs | 210 ---- .../plan_execution_dual_input_tests.rs | 312 ----- .../plan_execution_temporal_tests.rs | 604 ---------- .../tests/datafusion/plan_execution_tests.rs | 590 ---------- .../dispatch_arithmetic_tests.rs | 0 asap-query-engine/src/tests/mod.rs | 4 +- .../src/tests/native_range_query_tests.rs | 2 +- .../range_query_arithmetic_tests.rs | 0 .../structural_matching_tests.rs | 0 .../tests/test_utilities/engine_factories.rs | 72 -- 27 files changed, 4 insertions(+), 6568 deletions(-) delete mode 100644 asap-query-engine/src/bin/show_logical_plans.rs delete mode 100644 asap-query-engine/src/engines/logical/mod.rs delete mode 100644 asap-query-engine/src/engines/logical/plan_builder.rs delete mode 100644 asap-query-engine/src/engines/physical/accumulator_serde.rs delete mode 100644 asap-query-engine/src/engines/physical/conversion.rs delete mode 100644 asap-query-engine/src/engines/physical/mod.rs delete mode 100644 asap-query-engine/src/engines/physical/planner.rs delete mode 100644 asap-query-engine/src/engines/physical/precomputed_summary_read_exec.rs delete mode 100644 asap-query-engine/src/engines/physical/summary_infer_exec.rs delete mode 100644 asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs delete mode 100644 asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/mod.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_execution_tests.rs rename asap-query-engine/src/tests/{datafusion => }/dispatch_arithmetic_tests.rs (100%) rename asap-query-engine/src/tests/{datafusion => }/range_query_arithmetic_tests.rs (100%) rename asap-query-engine/src/tests/{datafusion => }/structural_matching_tests.rs (100%) diff --git a/Cargo.lock b/Cargo.lock index 2589722b..cdf6d875 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3403,7 +3403,6 @@ version = "0.5.0" dependencies = [ "anyhow", "arc-swap", - "arrow", "asap_planner", "asap_sketchlib", "asap_types", @@ -3416,8 +3415,6 @@ dependencies = [ "criterion", "csv", "dashmap 5.5.3", - "datafusion", - "datafusion_summary_library", "elastic_dsl_utilities", "figment", "flate2", diff --git a/asap-query-engine/Cargo.toml b/asap-query-engine/Cargo.toml index c31618a4..d780551d 100644 --- a/asap-query-engine/Cargo.toml +++ b/asap-query-engine/Cargo.toml @@ -8,7 +8,6 @@ edition.workspace = true promql_utilities.workspace = true sql_utilities.workspace = true asap_types.workspace = true -datafusion_summary_library.workspace = true asap_planner.workspace = true # Shared external (workspace) @@ -41,8 +40,6 @@ xxhash-rust = { version = "0.8", features = ["xxh32", "xxh64"] } base64 = "0.21" hex = "0.4" sqlparser = "0.59.0" -datafusion = "43" -arrow = "53.4.1" futures = "0.3" prost = "0.13" opentelemetry-proto = { version = "0.28", features = ["gen-tonic", "gen-tonic-messages", "metrics"] } diff --git a/asap-query-engine/src/bin/show_logical_plans.rs b/asap-query-engine/src/bin/show_logical_plans.rs deleted file mode 100644 index 39bd02b2..00000000 --- a/asap-query-engine/src/bin/show_logical_plans.rs +++ /dev/null @@ -1,528 +0,0 @@ -//! Standalone binary that constructs diverse QueryExecutionContext structures, -//! converts each to a DataFusion logical plan, and prints each plan along with -//! the schema of every edge and key internal variables. -//! -//! Covers 4 queries x multiple accumulator configurations = 10 test cases: -//! -//! 1. sum by (host) (data) — spatial sum -//! 2. quantile by (host) (0.5, data) — spatial quantile -//! 3. sum_over_time(data[1m]) — temporal sum -//! 4. quantile_over_time(0.5, data[1m]) — temporal quantile -//! -//! data has columns: host, service, region - -use datafusion::logical_expr::LogicalPlan; -use datafusion_summary_library::{PrecomputedSummaryRead, SummaryInfer, SummaryMergeMultiple}; -use promql_utilities::data_model::KeyByLabelNames; -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, -}; -use std::collections::HashMap; - -// ============================================================================ -// Context builders -// ============================================================================ - -/// Build a QueryExecutionContext with full control over all parameters. -#[allow(clippy::too_many_arguments)] -fn build_context( - metric: &str, - statistic: Statistic, - query_output_labels: Vec<&str>, - grouping_labels: Vec<&str>, - aggregated_labels: Vec<&str>, - agg_type_value: AggregationType, - agg_type_key: AggregationType, - agg_id_value: u64, - agg_id_key: u64, - keys_query: Option, - do_merge: bool, - is_exact_query: bool, - kwargs: HashMap, -) -> QueryExecutionContext { - QueryExecutionContext { - metric: metric.to_string(), - metadata: QueryMetadata { - query_output_labels: KeyByLabelNames { - labels: query_output_labels.into_iter().map(String::from).collect(), - }, - statistic_to_compute: statistic, - query_kwargs: kwargs, - }, - store_plan: StoreQueryPlan { - values_query: StoreQueryParams { - metric: metric.to_string(), - aggregation_id: agg_id_value, - start_timestamp: if do_merge { 1000 } else { 2000 }, - end_timestamp: 2000, - is_exact_query, - }, - keys_query, - }, - agg_info: AggregationIdInfo { - aggregation_id_for_key: agg_id_key, - aggregation_id_for_value: agg_id_value, - aggregation_type_for_key: agg_type_key, - aggregation_type_for_value: agg_type_value, - }, - do_merge, - spatial_filter: String::new(), - query_time: 2000, - grouping_labels: KeyByLabelNames { - labels: grouping_labels.into_iter().map(String::from).collect(), - }, - aggregated_labels: KeyByLabelNames { - labels: aggregated_labels.into_iter().map(String::from).collect(), - }, - } -} - -fn make_keys_query(metric: &str, agg_id: u64) -> StoreQueryParams { - StoreQueryParams { - metric: metric.to_string(), - aggregation_id: agg_id, - start_timestamp: 0, // DeltaSetAggregator reads from beginning of time - end_timestamp: 2000, - is_exact_query: false, // keys are always range queries - } -} - -// ============================================================================ -// Plan printing utilities -// ============================================================================ - -/// Recursively print the plan tree with indentation, showing each node's -/// explain text and output schema. -fn print_plan_tree(plan: &LogicalPlan, indent: usize) { - let prefix = " ".repeat(indent); - let connector = if indent > 0 { "└─► " } else { "" }; - - match plan { - LogicalPlan::Extension(ext) => { - // Print node name and explain text - println!("{prefix}{connector}{}", ext.node.name()); - - // Print detailed properties by downcasting each node - print_node_details(plan, indent + 2); - - // Print output schema - let schema = ext.node.schema(); - print!("{} schema: [", prefix); - for (i, field) in schema.fields().iter().enumerate() { - if i > 0 { - print!(", "); - } - print!("{}:{}", field.name(), field.data_type()); - } - println!("]"); - - // Recurse into inputs - let inputs = ext.node.inputs(); - for (i, input) in inputs.iter().enumerate() { - if inputs.len() > 1 { - println!("{} input {}:", prefix, i); - } - print_plan_tree(input, indent + 2); - } - } - _ => { - println!("{prefix}{connector}Unknown: {:?}", plan); - } - } -} - -/// Print detailed properties of a plan node by downcasting. -fn print_node_details(plan: &LogicalPlan, indent: usize) { - let prefix = " ".repeat(indent); - if let LogicalPlan::Extension(ext) = plan { - if let Some(infer) = ext.node.as_any().downcast_ref::() { - println!( - "{prefix}operations: {:?}", - infer - .operations - .iter() - .map(|op| format!("{}", op)) - .collect::>() - ); - println!("{prefix}output_names: {:?}", infer.output_names); - println!("{prefix}group_key_columns: {:?}", infer.group_key_columns); - println!("{prefix}has_keys_input: {}", infer.keys_input.is_some()); - } else if let Some(merge) = ext.node.as_any().downcast_ref::() { - println!("{prefix}group_by: {:?}", merge.group_by()); - println!("{prefix}sketch_column: {:?}", merge.sketch_column()); - println!("{prefix}summary_type: {}", merge.summary_type()); - } else if let Some(read) = ext.node.as_any().downcast_ref::() { - println!("{prefix}metric: {:?}", read.metric()); - println!("{prefix}aggregation_id: {}", read.aggregation_id()); - println!( - "{prefix}range: [{}, {}]", - read.start_timestamp(), - read.end_timestamp() - ); - println!("{prefix}is_exact_query: {}", read.is_exact_query()); - println!("{prefix}summary_type: {}", read.summary_type()); - println!("{prefix}output_labels: {:?}", read.output_labels()); - } - } -} - -/// Print key internal variables about a QueryExecutionContext. -fn print_context_variables(ctx: &QueryExecutionContext) { - let has_separate_keys = ctx.store_plan.keys_query.is_some() - && ctx.agg_info.aggregation_id_for_key != ctx.agg_info.aggregation_id_for_value; - let has_aggregated_labels = !ctx.aggregated_labels.labels.is_empty(); - - println!(" Internal variables:"); - println!(" has_separate_keys (dual input): {}", has_separate_keys); - println!( - " has_aggregated_labels (multi-population): {}", - has_aggregated_labels - ); - println!(" do_merge (temporal): {}", ctx.do_merge); - println!( - " keys_included: {}", - has_separate_keys || has_aggregated_labels - ); - println!( - " value_agg: {} (id={})", - ctx.agg_info.aggregation_type_for_value, ctx.agg_info.aggregation_id_for_value - ); - println!( - " key_agg: {} (id={})", - ctx.agg_info.aggregation_type_for_key, ctx.agg_info.aggregation_id_for_key - ); - println!( - " query_output_labels: {:?}", - ctx.metadata.query_output_labels.labels - ); - println!(" grouping_labels: {:?}", ctx.grouping_labels.labels); - println!(" aggregated_labels: {:?}", ctx.aggregated_labels.labels); - println!(" statistic: {:?}", ctx.metadata.statistic_to_compute); - if !ctx.metadata.query_kwargs.is_empty() { - println!(" query_kwargs: {:?}", ctx.metadata.query_kwargs); - } -} - -// ============================================================================ -// Test case definitions -// ============================================================================ - -struct TestCase { - title: String, - query: String, - description: String, - context: QueryExecutionContext, -} - -fn build_all_test_cases() -> Vec { - let metric = "data"; - let mut cases = Vec::new(); - - // ======================================================================== - // Query 1: sum by (host) (data) - // ======================================================================== - - // Case 1a: SumAccumulator only - // Simple single-population. Store groups by host, one Sum per host. - cases.push(TestCase { - title: "sum by (host) — SumAccumulator".into(), - query: "sum by (host) (data)".into(), - description: - "Single-population exact sum. Store groups by [host], one scalar sum per group key." - .into(), - context: build_context( - metric, - Statistic::Sum, - 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 - false, // not temporal - true, // exact (sliding window) - HashMap::new(), - ), - }); - - // Case 1b: MultipleSumAccumulator only (self-keyed) - // The accumulator internally tracks sums for each host value. - // Store doesn't group by host; the accumulator maps host -> sum. - cases.push(TestCase { - title: "sum by (host) — MultipleSumAccumulator (self-keyed)".into(), - query: "sum by (host) (data)".into(), - description: "Self-keyed multi-population. Store groups by [] (no spatial grouping). \ - MultipleSumAccumulator internally maps host -> sum." - .into(), - context: build_context( - metric, - Statistic::Sum, - vec!["host"], // query_output_labels - vec![], // grouping_labels (no store grouping) - vec!["host"], // aggregated_labels (host tracked internally) - AggregationType::MultipleSum, - AggregationType::MultipleSum, // same type = single agg_id - 42, - 42, - None, - false, - true, - HashMap::new(), - ), - }); - - // Case 1c: CountMinSketch + DeltaSetAggregator (dual-input) - // CountMinSketch estimates frequency per key; DeltaSetAggregator enumerates keys. - cases.push(TestCase { - title: "sum by (host) — CountMinSketch + DeltaSetAggregator (dual-input)".into(), - query: "sum by (host) (data)".into(), - description: "Dual-input plan. CountMinSketch for value estimation per host key, \ - DeltaSetAggregator enumerates which hosts exist." - .into(), - context: build_context( - metric, - Statistic::Sum, - vec!["host"], - vec![], // grouping_labels (no store grouping) - vec!["host"], // aggregated_labels - AggregationType::CountMinSketch, - AggregationType::DeltaSetAggregator, - 42, - 99, // different agg_ids - Some(make_keys_query(metric, 99)), - false, - true, - HashMap::new(), - ), - }); - - // ======================================================================== - // Query 2: quantile by (host) (0.5, data) - // ======================================================================== - - let mut q_kwargs = HashMap::new(); - q_kwargs.insert("quantile".to_string(), "0.5".to_string()); - - // Case 2a: KLL only - cases.push(TestCase { - title: "quantile by (host) (0.5) — KLL".into(), - query: "quantile by (host) (0.5, data)".into(), - description: - "Single-population quantile. Store groups by [host], one KLL sketch per group key." - .into(), - context: build_context( - metric, - Statistic::Quantile, - vec!["host"], - vec!["host"], - vec![], - AggregationType::DatasketchesKLL, - AggregationType::DatasketchesKLL, - 42, - 42, - None, - false, - true, - q_kwargs.clone(), - ), - }); - - // Case 2b: HydraKLL + DeltaSetAggregator (dual-input) - cases.push(TestCase { - title: "quantile by (host) (0.5) — HydraKLL + DeltaSetAggregator (dual-input)".into(), - query: "quantile by (host) (0.5, data)".into(), - description: "Dual-input quantile. HydraKLL has per-host KLL sketches internally. \ - DeltaSetAggregator enumerates which hosts exist." - .into(), - context: build_context( - metric, - Statistic::Quantile, - vec!["host"], - vec![], // no store grouping - vec!["host"], // host tracked internally - AggregationType::HydraKLL, - AggregationType::DeltaSetAggregator, - 42, - 99, - Some(make_keys_query(metric, 99)), - false, - true, - q_kwargs.clone(), - ), - }); - - // ======================================================================== - // Query 3: sum_over_time(data[1m]) - // Temporal — all labels preserved, do_merge=true - // ======================================================================== - - // Case 3a: SumAccumulator only - cases.push(TestCase { - title: "sum_over_time(data[1m]) — SumAccumulator".into(), - query: "sum_over_time(data[1m])".into(), - description: "Temporal sum, single-population. All labels preserved. \ - do_merge=true to merge tumbling windows across the 1m range." - .into(), - context: build_context( - metric, - Statistic::Sum, - vec!["host", "service", "region"], - vec!["host", "service", "region"], - vec![], - AggregationType::Sum, - AggregationType::Sum, - 42, - 42, - None, - true, // temporal merge - false, // tumbling window (range query) - HashMap::new(), - ), - }); - - // Case 3b: MultipleSumAccumulator only (self-keyed) - cases.push(TestCase { - title: "sum_over_time(data[1m]) — MultipleSumAccumulator (self-keyed)".into(), - query: "sum_over_time(data[1m])".into(), - description: "Temporal sum, self-keyed multi-population. Store groups by [host]. \ - MultipleSumAccumulator internally maps (service, region) -> sum." - .into(), - context: build_context( - metric, - Statistic::Sum, - vec!["host", "service", "region"], - vec!["host"], // store groups by host only - vec!["service", "region"], // rest tracked internally - AggregationType::MultipleSum, - AggregationType::MultipleSum, - 42, - 42, - None, - true, - false, - HashMap::new(), - ), - }); - - // Case 3c: CountMinSketch + DeltaSetAggregator (dual-input) - cases.push(TestCase { - title: "sum_over_time(data[1m]) — CountMinSketch + DeltaSetAggregator (dual-input)".into(), - query: "sum_over_time(data[1m])".into(), - description: "Temporal sum, dual-input. Store groups by [host]. \ - CountMinSketch estimates per (service, region). \ - DeltaSetAggregator enumerates (service, region) keys." - .into(), - context: build_context( - metric, - Statistic::Sum, - vec!["host", "service", "region"], - vec!["host"], - vec!["service", "region"], - AggregationType::CountMinSketch, - AggregationType::DeltaSetAggregator, - 42, - 99, - Some(make_keys_query(metric, 99)), - true, - false, - HashMap::new(), - ), - }); - - // ======================================================================== - // Query 4: quantile_over_time(0.5, data[1m]) - // Temporal — all labels preserved, do_merge=true - // ======================================================================== - - // Case 4a: KLL only - cases.push(TestCase { - title: "quantile_over_time(0.5, data[1m]) — KLL".into(), - query: "quantile_over_time(0.5, data[1m])".into(), - description: "Temporal quantile, single-population. All labels preserved. \ - One KLL sketch per (host, service, region) group." - .into(), - context: build_context( - metric, - Statistic::Quantile, - vec!["host", "service", "region"], - vec!["host", "service", "region"], - vec![], - AggregationType::DatasketchesKLL, - AggregationType::DatasketchesKLL, - 42, - 42, - None, - true, - false, - q_kwargs.clone(), - ), - }); - - // Case 4b: HydraKLL + DeltaSetAggregator (dual-input) - cases.push(TestCase { - title: "quantile_over_time(0.5, data[1m]) — HydraKLL + DeltaSetAggregator (dual-input)" - .into(), - query: "quantile_over_time(0.5, data[1m])".into(), - description: "Temporal quantile, dual-input. Store groups by [host]. \ - HydraKLL has per-(service, region) KLL sketches. \ - DeltaSetAggregator enumerates (service, region) keys." - .into(), - context: build_context( - metric, - Statistic::Quantile, - vec!["host", "service", "region"], - vec!["host"], - vec!["service", "region"], - AggregationType::HydraKLL, - AggregationType::DeltaSetAggregator, - 42, - 99, - Some(make_keys_query(metric, 99)), - true, - false, - q_kwargs.clone(), - ), - }); - - cases -} - -// ============================================================================ -// Main -// ============================================================================ - -fn main() { - let cases = build_all_test_cases(); - - for (i, case) in cases.iter().enumerate() { - println!("╔══════════════════════════════════════════════════════════════════════"); - println!("║ Case {}: {}", i + 1, case.title); - println!("║ Query: {}", case.query); - println!("║ {}", case.description); - println!("╚══════════════════════════════════════════════════════════════════════"); - println!(); - - // Print key internal variables - print_context_variables(&case.context); - println!(); - - // Convert to logical plan - match case.context.to_logical_plan() { - Ok(plan) => { - println!(" Logical Plan Tree:"); - println!(" ──────────────────"); - print_plan_tree(&plan, 2); - } - Err(e) => { - println!(" ERROR converting to logical plan: {}", e); - } - } - - println!(); - println!(); - } -} diff --git a/asap-query-engine/src/engines/logical/mod.rs b/asap-query-engine/src/engines/logical/mod.rs deleted file mode 100644 index 435ca52b..00000000 --- a/asap-query-engine/src/engines/logical/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -//! Logical Plan Builders for DataFusion -//! -//! This module converts QueryExecutionContext into DataFusion LogicalPlan trees -//! using the custom extension nodes defined in datafusion_summary_library. - -pub mod plan_builder; diff --git a/asap-query-engine/src/engines/logical/plan_builder.rs b/asap-query-engine/src/engines/logical/plan_builder.rs deleted file mode 100644 index 813f0241..00000000 --- a/asap-query-engine/src/engines/logical/plan_builder.rs +++ /dev/null @@ -1,1014 +0,0 @@ -//! Plan Builder Module -//! -//! Converts QueryExecutionContext to DataFusion LogicalPlan for OnlySpatial queries. -//! This enables plan-based execution as an alternative to the existing pipeline. - -use arrow::datatypes::{DataType, Field}; -use datafusion::common::{DFSchema, DFSchemaRef}; -use datafusion::error::DataFusionError; -use datafusion::logical_expr::{ - binary_expr, Expr as DFExpr, Extension, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, - SubqueryAlias, -}; -use datafusion::prelude::{col, lit}; -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::{AggregationType, Statistic}; -use std::sync::Arc; - -use crate::engines::simple_engine::{QueryExecutionContext, StoreQueryParams}; - -/// Extension trait for building DataFusion logical plans from QueryExecutionContext -impl QueryExecutionContext { - /// Convert this execution context to a DataFusion LogicalPlan. - /// - /// The resulting plan structure for single-population queries: - /// ```text - /// SummaryInfer (extract values from summaries) - /// └─► SummaryMergeMultiple (merge summaries with same group key) - /// └─► PrecomputedSummaryRead (read from store) - /// ``` - /// - /// For multi-population queries (separate keys_query): - /// ```text - /// SummaryInfer (dual-input: value sketch + keys enumeration) - /// ├─► input 0: SummaryMergeMultiple (values) - /// │ └─► PrecomputedSummaryRead (values agg_id) - /// └─► input 1: SummaryMergeMultiple (keys) - /// └─► PrecomputedSummaryRead (keys agg_id) - /// ``` - pub fn to_logical_plan(&self) -> Result { - let has_separate_keys = self.store_plan.keys_query.is_some() - && self.agg_info.aggregation_id_for_key != self.agg_info.aggregation_id_for_value; - - // 1. Map aggregation type to SummaryType (SketchType) for values - let summary_type = self.map_aggregation_type_to_summary_type()?; - - // Determine labels for the values branch (store read/merge). - // For multi-population (dual-input or self-keyed): use grouping_labels (store GROUP BY) - // For single-population: use query_output_labels - let has_aggregated_labels = !self.aggregated_labels.labels.is_empty(); - let values_labels = if has_separate_keys || has_aggregated_labels { - self.grouping_labels.labels.to_vec() - } else { - self.get_output_label_names() - }; - - // Sub-key labels come from aggregated_labels (labels that key the accumulator internally) - let sub_key_labels: Vec = self.aggregated_labels.labels.to_vec(); - - // 2. Build values branch: Read -> Merge - let values_merge_plan = self.build_read_merge_branch( - &self.store_plan.values_query, - &values_labels, - &summary_type, - )?; - - // 3. Map statistic to InferOperation - let infer_operation = self.map_statistic_to_infer_operation()?; - - if has_separate_keys { - let keys_query = self.store_plan.keys_query.as_ref().unwrap(); - - // Map keys aggregation type to SketchType - let keys_summary_type = self.map_key_aggregation_type_to_summary_type()?; - - // Build keys branch: Read -> Merge (using same spatial labels) - let keys_merge_plan = - self.build_read_merge_branch(keys_query, &values_labels, &keys_summary_type)?; - - // Create dual-input SummaryInfer - let infer = SummaryInfer::new( - Arc::new(values_merge_plan), - vec![infer_operation], - vec!["value".to_string()], - ) - .map_err(|e| DataFusionError::Plan(format!("Failed to create SummaryInfer: {}", e)))? - .with_keys_input(Arc::new(keys_merge_plan)) - .with_group_key_columns(sub_key_labels, None) - .map_err(|e| { - DataFusionError::Plan(format!("Failed to set group_key_columns: {}", e)) - })?; - - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(infer), - })) - } else { - // Single-input path - let mut infer = SummaryInfer::new( - Arc::new(values_merge_plan), - vec![infer_operation], - vec!["value".to_string()], - ) - .map_err(|e| DataFusionError::Plan(format!("Failed to create SummaryInfer: {}", e)))?; - - if !sub_key_labels.is_empty() { - // Self-keyed multi-pop: set sub-key columns so the output schema - // includes them and the physical operator knows to enumerate keys. - infer = infer - .with_group_key_columns(sub_key_labels, None) - .map_err(|e| { - DataFusionError::Plan(format!("Failed to set group_key_columns: {}", e)) - })?; - } - - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(infer), - })) - } - } - - /// Build a Read -> Merge branch for a given store query. - fn build_read_merge_branch( - &self, - query_params: &StoreQueryParams, - labels: &[String], - summary_type: &SketchType, - ) -> Result { - let read_schema = self.build_read_schema(labels)?; - - let read = PrecomputedSummaryRead::new( - self.metric.clone(), - query_params.aggregation_id, - query_params.start_timestamp, - query_params.end_timestamp, - query_params.is_exact_query, - labels.to_vec(), - summary_type.clone(), - read_schema, - ); - let read_plan = LogicalPlan::Extension(Extension { - node: Arc::new(read), - }); - - let merge = SummaryMergeMultiple::new( - Arc::new(read_plan), - labels.to_vec(), - "sketch".to_string(), - summary_type.clone(), - ); - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(merge), - })) - } - - /// Get output label names from the query metadata - fn get_output_label_names(&self) -> Vec { - self.metadata.query_output_labels.labels.to_vec() - } - - /// Build schema for PrecomputedSummaryRead: [label columns, sketch column] - fn build_read_schema(&self, output_labels: &[String]) -> Result { - let mut fields: Vec<(Option, Arc)> = Vec::new(); - - // Add label columns (Utf8, nullable) - for label in output_labels { - fields.push((None, Arc::new(Field::new(label, DataType::Utf8, true)))); - } - - // Add sketch column (Binary, not nullable) - fields.push(( - None, - Arc::new(Field::new("sketch", DataType::Binary, false)), - )); - - let schema = DFSchema::new_with_metadata(fields, Default::default()) - .map_err(|e| DataFusionError::Plan(format!("Failed to create read schema: {}", e)))?; - - Ok(Arc::new(schema)) - } - - /// Map Statistic enum to InferOperation - pub(crate) fn map_statistic_to_infer_operation( - &self, - ) -> Result { - match self.metadata.statistic_to_compute { - Statistic::Sum => Ok(InferOperation::ExtractSum), - Statistic::Min => Ok(InferOperation::ExtractMin), - Statistic::Max => Ok(InferOperation::ExtractMax), - Statistic::Count => Ok(InferOperation::ExtractCount), - Statistic::Increase => Ok(InferOperation::ExtractIncrease), - Statistic::Rate => Ok(InferOperation::ExtractRate), - Statistic::Quantile => { - // Extract quantile parameter from query_kwargs - let q = self - .metadata - .query_kwargs - .get("quantile") - .and_then(|s| s.parse::().ok()) - .unwrap_or(0.5); - Ok(InferOperation::quantile(q)) - } - Statistic::Cardinality => Ok(InferOperation::CountDistinct), - Statistic::Topk => { - // Extract k parameter from query_kwargs - let k = self - .metadata - .query_kwargs - .get("k") - .and_then(|s| s.parse::().ok()) - .unwrap_or(10); - Ok(InferOperation::TopK(k)) - } - } - } - - /// Map aggregation type to SketchType (SummaryType) for the value accumulator - fn map_aggregation_type_to_summary_type(&self) -> Result { - Self::agg_type_to_sketch_type(self.agg_info.aggregation_type_for_value) - } - - /// Map aggregation type to SketchType (SummaryType) for the key accumulator - fn map_key_aggregation_type_to_summary_type(&self) -> Result { - Self::agg_type_to_sketch_type(self.agg_info.aggregation_type_for_key) - } - - fn agg_type_to_sketch_type(agg_type: AggregationType) -> Result { - match agg_type { - 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:?}" - ))), - } - } -} - -// ============================================================================ -// Binary arithmetic plan builders (standalone functions, not on impl block) -// ============================================================================ - -/// Map a PromQL `TokenType` to the corresponding DataFusion `Operator`. -/// Returns an error for non-arithmetic operators. -pub fn token_type_to_df_operator(op: &token::TokenType) -> Result { - match op.id() { - id if id == T_ADD => Ok(Operator::Plus), - id if id == T_SUB => Ok(Operator::Minus), - id if id == T_MUL => Ok(Operator::Multiply), - id if id == T_DIV => Ok(Operator::Divide), - id if id == T_MOD => Ok(Operator::Modulo), - id if id == T_POW => Ok(Operator::BitwiseXor), // Note: bitwise XOR used as proxy for ^ - _ => Err(DataFusionError::Plan(format!( - "Unsupported binary operator for arithmetic plan: {}", - op - ))), - } -} - -/// Builds a DataFusion logical plan for a vector op vector binary expression: -/// -/// ```text -/// Projection(lhs.label1 AS label1, ..., lhs.value OP rhs.value AS value) -/// └── Join(inner, on = label_columns) -/// ├── SubqueryAlias("lhs") └── lhs_plan -/// └── SubqueryAlias("rhs") └── rhs_plan -/// ``` -/// -/// The `label_columns` are the label names shared by both sides; the join is -/// an inner join on those columns. -pub fn build_binary_vector_plan( - lhs_plan: LogicalPlan, - rhs_plan: LogicalPlan, - op: &token::TokenType, - label_columns: Vec, -) -> Result { - let df_op = token_type_to_df_operator(op)?; - - // Wrap each side in a SubqueryAlias so columns are qualified (lhs.x, rhs.x) - let lhs_aliased = SubqueryAlias::try_new(Arc::new(lhs_plan), "lhs")?; - let rhs_aliased = SubqueryAlias::try_new(Arc::new(rhs_plan), "rhs")?; - - // Build the join keys: qualified column names for each label column. - // The `join` function expects Vec> (i.e. qualified col names). - let join_keys_left: Vec = label_columns.iter().map(|c| format!("lhs.{}", c)).collect(); - let join_keys_right: Vec = label_columns.iter().map(|c| format!("rhs.{}", c)).collect(); - - let joined_plan = LogicalPlanBuilder::from(LogicalPlan::SubqueryAlias(lhs_aliased)) - .join( - LogicalPlan::SubqueryAlias(rhs_aliased), - JoinType::Inner, - (join_keys_left, join_keys_right), - None, - )? - .build()?; - - // Projection: pass through label columns from lhs, compute value = lhs.value OP rhs.value - let mut proj_exprs: Vec = label_columns - .iter() - .map(|c| col(format!("lhs.{}", c)).alias(c.as_str())) - .collect(); - let value_expr = binary_expr(col("lhs.value"), df_op, col("rhs.value")).alias("value"); - proj_exprs.push(value_expr); - - LogicalPlanBuilder::from(joined_plan) - .project(proj_exprs)? - .build() -} - -/// Builds a DataFusion logical plan for a scalar op vector (or vector op scalar) expression: -/// -/// ```text -/// Projection(label1, ..., scalar OP value AS value) -/// └── vector_plan -/// ``` -/// -/// If `scalar_on_left` is true the expression is `scalar OP value`; -/// otherwise it is `value OP scalar`. -pub fn build_scalar_plan( - vector_plan: LogicalPlan, - scalar: f64, - op: &token::TokenType, - scalar_on_left: bool, - label_columns: Vec, -) -> Result { - let df_op = token_type_to_df_operator(op)?; - - let value_expr = if scalar_on_left { - binary_expr(lit(scalar), df_op, col("value")).alias("value") - } else { - binary_expr(col("value"), df_op, lit(scalar)).alias("value") - }; - - let mut proj_exprs: Vec = label_columns.iter().map(|c| col(c.as_str())).collect(); - proj_exprs.push(value_expr); - - LogicalPlanBuilder::from(vector_plan) - .project(proj_exprs)? - .build() -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::data_model::AggregationIdInfo; - use crate::engines::simple_engine::{QueryMetadata, StoreQueryParams, StoreQueryPlan}; - use promql_utilities::data_model::KeyByLabelNames; - - use std::collections::HashMap; - - fn create_test_context( - metric: &str, - statistic: Statistic, - output_labels: Vec<&str>, - aggregation_type: AggregationType, - ) -> QueryExecutionContext { - 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: AggregationType, - kwargs: HashMap, - ) -> QueryExecutionContext { - let mut ctx = create_test_context_with_keys( - metric, - statistic, - output_labels, - aggregation_type, - None, - aggregation_type, - ); - ctx.metadata.query_kwargs = kwargs; - ctx - } - - fn create_test_context_with_keys( - metric: &str, - statistic: Statistic, - output_labels: Vec<&str>, - aggregation_type_for_value: AggregationType, - keys_query: Option, - aggregation_type_for_key: AggregationType, - ) -> QueryExecutionContext { - // Default: grouping_labels == query_output_labels - create_test_context_with_keys_and_grouping( - metric, - statistic, - output_labels.clone(), - aggregation_type_for_value, - keys_query, - aggregation_type_for_key, - output_labels, - ) - } - - fn create_test_context_with_keys_and_grouping( - metric: &str, - statistic: Statistic, - output_labels: Vec<&str>, - aggregation_type_for_value: AggregationType, - keys_query: Option, - aggregation_type_for_key: AggregationType, - grouping_label_strs: Vec<&str>, - ) -> QueryExecutionContext { - // Default: aggregated_labels = output_labels - grouping_labels - let aggregated: Vec<&str> = output_labels - .iter() - .filter(|l| !grouping_label_strs.contains(l)) - .copied() - .collect(); - create_test_context_full( - metric, - statistic, - output_labels, - aggregation_type_for_value, - keys_query, - aggregation_type_for_key, - grouping_label_strs, - aggregated, - ) - } - - #[allow(clippy::too_many_arguments)] - fn create_test_context_full( - metric: &str, - statistic: Statistic, - output_labels: Vec<&str>, - aggregation_type_for_value: AggregationType, - keys_query: Option, - aggregation_type_for_key: AggregationType, - grouping_label_strs: Vec<&str>, - aggregated_label_strs: Vec<&str>, - ) -> QueryExecutionContext { - let aggregation_id_for_key = match &keys_query { - Some(kq) => kq.aggregation_id, - None => 42, // same as value when no separate keys - }; - let output_labels_vec: Vec = output_labels.into_iter().map(String::from).collect(); - let query_output_labels = KeyByLabelNames { - labels: output_labels_vec, - }; - let grouping_labels = KeyByLabelNames { - labels: grouping_label_strs.into_iter().map(String::from).collect(), - }; - let aggregated_labels = KeyByLabelNames { - labels: aggregated_label_strs - .into_iter() - .map(String::from) - .collect(), - }; - QueryExecutionContext { - metric: metric.to_string(), - metadata: QueryMetadata { - query_output_labels, - statistic_to_compute: statistic, - query_kwargs: HashMap::new(), - }, - store_plan: StoreQueryPlan { - values_query: StoreQueryParams { - metric: metric.to_string(), - aggregation_id: 42, - start_timestamp: 1000, - end_timestamp: 2000, - is_exact_query: true, - }, - keys_query, - }, - agg_info: AggregationIdInfo { - aggregation_id_for_key, - aggregation_id_for_value: 42, - aggregation_type_for_key, - aggregation_type_for_value, - }, - do_merge: false, - spatial_filter: String::new(), - query_time: 2000, - grouping_labels, - aggregated_labels, - } - } - - #[test] - fn test_to_logical_plan_creates_correct_structure() { - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::Sum, - ); - - let plan = context.to_logical_plan().unwrap(); - - // Root should be SummaryInfer - match &plan { - LogicalPlan::Extension(ext) => { - assert_eq!(ext.node.name(), "SummaryInfer"); - } - _ => panic!("Expected Extension node"), - } - } - - #[test] - fn test_to_logical_plan_with_multiple_labels() { - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host", "region", "service"], - AggregationType::Sum, - ); - - let plan = context.to_logical_plan().unwrap(); - - // Verify plan was created successfully - assert!(matches!(plan, LogicalPlan::Extension(_))); - } - - #[test] - fn test_map_statistic_to_infer_operation() { - 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"], - AggregationType::MinMax, - ); - assert!(matches!( - context.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractMin - )); - - let context = create_test_context( - "test", - Statistic::Max, - vec!["host"], - AggregationType::MinMax, - ); - assert!(matches!( - context.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractMax - )); - } - - #[test] - fn test_map_aggregation_type_to_summary_type() { - let context = - create_test_context("test", Statistic::Sum, vec!["host"], AggregationType::Sum); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::Sum - ); - - let context = create_test_context( - "test", - Statistic::Increase, - vec!["host"], - AggregationType::Increase, - ); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::Increase - ); - - let context = create_test_context( - "test", - Statistic::Quantile, - vec!["host"], - AggregationType::DatasketchesKLL, - ); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::KLL - ); - } - - // ======================================================================== - // Helper to walk the plan tree and collect node names top-down - // ======================================================================== - - fn collect_plan_node_names(plan: &LogicalPlan) -> Vec { - let mut names = Vec::new(); - collect_plan_node_names_recursive(plan, &mut names); - names - } - - fn collect_plan_node_names_recursive(plan: &LogicalPlan, names: &mut Vec) { - match plan { - LogicalPlan::Extension(ext) => { - names.push(ext.node.name().to_string()); - for input in ext.node.inputs() { - collect_plan_node_names_recursive(input, names); - } - } - _ => { - names.push( - format!("{:?}", plan) - .split('(') - .next() - .unwrap_or("Unknown") - .to_string(), - ); - } - } - } - - /// Helper to extract the SummaryMergeMultiple node from a plan tree - fn extract_merge_node(plan: &LogicalPlan) -> Option<&SummaryMergeMultiple> { - match plan { - LogicalPlan::Extension(ext) => { - if let Some(merge) = ext.node.as_any().downcast_ref::() { - return Some(merge); - } - for input in ext.node.inputs() { - if let Some(merge) = extract_merge_node(input) { - return Some(merge); - } - } - None - } - _ => None, - } - } - - /// Helper to extract the PrecomputedSummaryRead node from a plan tree - fn extract_read_node(plan: &LogicalPlan) -> Option<&PrecomputedSummaryRead> { - match plan { - LogicalPlan::Extension(ext) => { - if let Some(read) = ext.node.as_any().downcast_ref::() { - return Some(read); - } - for input in ext.node.inputs() { - if let Some(read) = extract_read_node(input) { - return Some(read); - } - } - None - } - _ => None, - } - } - - /// Helper to extract the SummaryInfer node from the root of a plan - fn extract_infer_node(plan: &LogicalPlan) -> Option<&SummaryInfer> { - match plan { - LogicalPlan::Extension(ext) => ext.node.as_any().downcast_ref::(), - _ => None, - } - } - - /// Count PrecomputedSummaryRead nodes in the plan tree - fn count_read_nodes(plan: &LogicalPlan) -> usize { - let names = collect_plan_node_names(plan); - names - .iter() - .filter(|n| *n == "PrecomputedSummaryRead") - .count() - } - - // ======================================================================== - // MultipleSumAccumulator (HydraSum) tests - // ======================================================================== - - #[test] - fn test_multiple_sum_accumulator_maps_to_hydra_sum() { - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::MultipleSum, - ); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::MultipleSum - ); - } - - #[test] - fn test_multiple_sum_accumulator_plan_builds() { - // MultipleSumAccumulator is a Hydra (multi-population) accumulator. - // The current plan only builds a single-population SummaryInfer, which - // won't correctly query sub-populations at execution time. - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::MultipleSum, - ); - - let plan = context.to_logical_plan().unwrap(); - let node_names = collect_plan_node_names(&plan); - assert_eq!( - node_names, - vec![ - "SummaryInfer", - "SummaryMergeMultiple", - "PrecomputedSummaryRead" - ] - ); - - // Verify the summary type propagates correctly through the plan - let merge = extract_merge_node(&plan).expect("Should have a SummaryMergeMultiple node"); - assert_eq!(*merge.summary_type(), SketchType::MultipleSum); - - let read = extract_read_node(&plan).expect("Should have a PrecomputedSummaryRead node"); - assert_eq!(*read.summary_type(), SketchType::MultipleSum); - } - - #[test] - fn test_multiple_sum_accumulator_single_pop_no_subkeys() { - // MultipleSumAccumulator without a separate keys_query stays single-population. - // No sub-key columns because there's no keys branch to enumerate from. - // To properly query Hydra types, a keys_query with a DeltaSetAggregator - // should be provided — see test_delta_set_aggregator_dual_input_plan. - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::MultipleSum, - ); - - let plan = context.to_logical_plan().unwrap(); - - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert!( - infer.group_key_columns.is_empty(), - "Single-pop (no keys_query): group_key_columns should be empty" - ); - assert!( - infer.keys_input.is_none(), - "Single-pop: should not have keys_input" - ); - } - - // ======================================================================== - // CountMinSketch for values_plan tests - // ======================================================================== - - #[test] - fn test_count_min_sketch_maps_correctly() { - let context = create_test_context( - "http_requests", - Statistic::Count, - vec!["host"], - AggregationType::CountMinSketch, - ); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::CountMinSketch - ); - } - - #[test] - fn test_count_min_sketch_plan_builds() { - // CountMinSketch is a multi-population frequency sketch. - // Like Hydra types, it requires a sub-key to query (FrequencyEstimate, etc.) - let context = create_test_context( - "http_requests", - Statistic::Count, - vec!["host"], - AggregationType::CountMinSketch, - ); - - let plan = context.to_logical_plan().unwrap(); - let node_names = collect_plan_node_names(&plan); - assert_eq!( - node_names, - vec![ - "SummaryInfer", - "SummaryMergeMultiple", - "PrecomputedSummaryRead" - ] - ); - - // Verify summary type - let merge = extract_merge_node(&plan).expect("Should have SummaryMergeMultiple"); - assert_eq!(*merge.summary_type(), SketchType::CountMinSketch); - } - - #[test] - fn test_count_min_sketch_single_pop_no_subkeys() { - // CountMinSketch without a separate keys_query stays single-population. - // Same as MultipleSumAccumulator — need a keys_query for dual-input. - let context = create_test_context( - "http_requests", - Statistic::Count, - vec!["host"], - AggregationType::CountMinSketch, - ); - - let plan = context.to_logical_plan().unwrap(); - - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert!(infer.group_key_columns.is_empty()); - assert!(infer.keys_input.is_none()); - } - - // ======================================================================== - // DeltaSetAggregator for keys_plan tests - // ======================================================================== - - #[test] - fn test_delta_set_aggregator_dual_input_plan() { - // When keys_query is set with a different agg_id (DeltaSetAggregator for key - // enumeration), to_logical_plan() builds a dual-input SummaryInfer with both - // a values branch and a keys branch. - let keys_query = StoreQueryParams { - metric: "http_requests".to_string(), - aggregation_id: 99, // Different agg ID for keys - start_timestamp: 0, - end_timestamp: 2000, - is_exact_query: false, - }; - let context = create_test_context_with_keys( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::MultipleSum, // values use Hydra - Some(keys_query), - AggregationType::DeltaSetAggregator, // keys use DeltaSet - ); - - let plan = context.to_logical_plan().unwrap(); - - // The plan tree should now have 5 nodes: SummaryInfer with 2 branches - let node_names = collect_plan_node_names(&plan); - assert_eq!( - node_names, - vec![ - "SummaryInfer", - "SummaryMergeMultiple", // values branch - "PrecomputedSummaryRead", // values read - "SummaryMergeMultiple", // keys branch - "PrecomputedSummaryRead", // keys read - ] - ); - - // Verify there are 2 PrecomputedSummaryRead nodes - assert_eq!(count_read_nodes(&plan), 2); - - // The SummaryInfer should have a keys_input - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert!( - infer.keys_input.is_some(), - "SummaryInfer should have keys_input" - ); - } - - // ======================================================================== - // HydraKLL for values_plan + DeltaSetAggregator for keys_plan - // ======================================================================== - - #[test] - fn test_hydra_kll_with_delta_set_keys_plan_builds() { - // HydraKLL for quantile queries with DeltaSetAggregator for key enumeration. - // This is a realistic configuration: HydraKLL stores per-key quantile sketches, - // and DeltaSetAggregator tracks which keys exist. - // grouping_labels = ["host"], sub-keys = ["endpoint"] - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.95".to_string()); - - let keys_query = StoreQueryParams { - metric: "request_duration".to_string(), - aggregation_id: 99, - start_timestamp: 0, - end_timestamp: 2000, - is_exact_query: false, - }; - let mut context = create_test_context_with_keys_and_grouping( - "request_duration", - Statistic::Quantile, - vec!["host", "endpoint"], // query_output_labels - AggregationType::HydraKLL, - Some(keys_query), - AggregationType::DeltaSetAggregator, - vec!["host"], // grouping_labels (spatial) - ); - context.metadata.query_kwargs = kwargs; - - let plan = context.to_logical_plan().unwrap(); - - // Verify dual-input plan structure (5 nodes) - let node_names = collect_plan_node_names(&plan); - assert_eq!( - node_names, - vec![ - "SummaryInfer", - "SummaryMergeMultiple", // values branch - "PrecomputedSummaryRead", // values read - "SummaryMergeMultiple", // keys branch - "PrecomputedSummaryRead", // keys read - ] - ); - - // Verify types propagate in the values branch - let merge = extract_merge_node(&plan).expect("Should have SummaryMergeMultiple"); - assert_eq!(*merge.summary_type(), SketchType::HydraKLL); - - let read = extract_read_node(&plan).expect("Should have PrecomputedSummaryRead"); - assert_eq!(*read.summary_type(), SketchType::HydraKLL); - // Values branch uses grouping_labels, not query_output_labels - assert_eq!(read.output_labels(), &["host"]); - - // Verify SummaryInfer has sub-key columns - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert_eq!(infer.group_key_columns, vec!["endpoint"]); - } - - #[test] - fn test_hydra_kll_with_delta_set_keys_dual_input_with_subkeys() { - // HydraKLL with DeltaSetAggregator keys, with sub-key columns. - // output_labels = ["host", "endpoint"], grouping_labels = ["host"] - // => sub_key_labels = ["endpoint"] - let keys_query = StoreQueryParams { - metric: "request_duration".to_string(), - aggregation_id: 99, - start_timestamp: 0, - end_timestamp: 2000, - is_exact_query: false, - }; - let context = create_test_context_with_keys_and_grouping( - "request_duration", - Statistic::Quantile, - vec!["host", "endpoint"], // query_output_labels - AggregationType::HydraKLL, - Some(keys_query), - AggregationType::DeltaSetAggregator, - vec!["host"], // grouping_labels (spatial store labels) - ); - - let plan = context.to_logical_plan().unwrap(); - - // Plan should have 2 PrecomputedSummaryRead nodes - assert_eq!(count_read_nodes(&plan), 2); - - // SummaryInfer should have keys_input and group_key_columns = ["endpoint"] - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert!(infer.keys_input.is_some(), "Should have keys_input"); - assert_eq!( - infer.group_key_columns, - vec!["endpoint"], - "Sub-key columns should be query_output_labels minus grouping_labels" - ); - } - - // ======================================================================== - // Quantile kwargs propagation test - // ======================================================================== - - #[test] - fn test_quantile_kwargs_propagate_to_infer_operation() { - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.99".to_string()); - let context = create_test_context_with_kwargs( - "latency", - Statistic::Quantile, - vec!["host"], - AggregationType::DatasketchesKLL, - kwargs, - ); - - match context.map_statistic_to_infer_operation().unwrap() { - InferOperation::Quantile(q) => { - // 0.99 * 10000 = 9900 - assert_eq!(q, 9900, "Expected q=9900 (0.99), got {}", q); - } - other => panic!("Expected Quantile, got {:?}", other), - } - } - - #[test] - fn test_quantile_defaults_to_median_when_no_kwargs() { - let context = create_test_context( - "latency", - Statistic::Quantile, - vec!["host"], - AggregationType::DatasketchesKLL, - ); - - match context.map_statistic_to_infer_operation().unwrap() { - InferOperation::Quantile(q) => { - // 0.5 * 10000 = 5000 - assert_eq!(q, 5000, "Expected q=5000 (0.5), got {}", q); - } - other => panic!("Expected Quantile, got {:?}", other), - } - } -} diff --git a/asap-query-engine/src/engines/mod.rs b/asap-query-engine/src/engines/mod.rs index 88b65a64..b81b5ac6 100644 --- a/asap-query-engine/src/engines/mod.rs +++ b/asap-query-engine/src/engines/mod.rs @@ -1,6 +1,4 @@ -pub mod logical; pub(crate) mod merge_utils; -pub mod physical; pub mod query_result; pub mod simple_engine; pub mod window_merger; diff --git a/asap-query-engine/src/engines/physical/accumulator_serde.rs b/asap-query-engine/src/engines/physical/accumulator_serde.rs deleted file mode 100644 index 6b7e480e..00000000 --- a/asap-query-engine/src/engines/physical/accumulator_serde.rs +++ /dev/null @@ -1,363 +0,0 @@ -//! Accumulator Serialization/Deserialization Registry -//! -//! Provides functions to deserialize bytes back to accumulator objects -//! based on the SummaryType (SketchType). -//! -//! Note: This module assumes accumulators are serialized using the Arroyo format -//! (MessagePack serialization via rmp-serde). - -use datafusion::error::DataFusionError; -use datafusion_summary_library::SketchType; - -use crate::data_model::{MultipleSubpopulationAggregate, SingleSubpopulationAggregate}; -use crate::precompute_operators::{ - CountMinSketchAccumulator, DatasketchesKLLAccumulator, DeltaSetAggregatorAccumulator, - HllAccumulator, HydraKllSketchAccumulator, MultipleIncreaseAccumulator, MultipleSumAccumulator, - SetAggregatorAccumulator, SumAccumulator, -}; -use crate::AggregateCore; - -/// Deserialize bytes to an accumulator based on the summary type. -/// -/// This function dispatches to the appropriate accumulator deserializer -/// based on the SummaryType enum. Expects Arroyo format (MessagePack). -/// -/// # Arguments -/// * `bytes` - Serialized accumulator data (Arroyo/MessagePack format) -/// * `summary_type` - Type of the summary (determines which deserializer to use) -/// -/// # Returns -/// A boxed AggregateCore trait object -pub fn deserialize_accumulator( - bytes: &[u8], - summary_type: &SketchType, -) -> Result, DataFusionError> { - match summary_type { - // Single-population exact aggregators - SketchType::Sum => { - let acc = SumAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize Sum: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::Increase => Err(DataFusionError::NotImplemented( - "Increase Arroyo deserialization not implemented".to_string(), - )), - SketchType::MinMax => Err(DataFusionError::NotImplemented( - "MinMax Arroyo deserialization not implemented".to_string(), - )), - - // Quantile sketches - SketchType::KLL => { - let acc = - DatasketchesKLLAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize KLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::HydraKLL => { - let acc = - HydraKllSketchAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize HydraKLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - - // Set aggregators - SketchType::SetAggregator => { - let acc = - SetAggregatorAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize SetAggregator: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::DeltaSetAggregator => { - let acc = DeltaSetAggregatorAccumulator::deserialize_from_bytes_arroyo(bytes).map_err( - |e| { - DataFusionError::Internal(format!( - "Failed to deserialize DeltaSetAggregator: {}", - e - )) - }, - )?; - Ok(Box::new(acc)) - } - - // Multi-population exact aggregators - SketchType::MultipleIncrease => { - let acc = - MultipleIncreaseAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize MultipleIncrease: {}", - e - )) - })?; - Ok(Box::new(acc)) - } - SketchType::MultipleSum => { - let acc = - MultipleSumAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize MultipleSum: {}", e)) - })?; - Ok(Box::new(acc)) - } - - // Frequency sketches - SketchType::CountMinSketch => { - let acc = - CountMinSketchAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize CountMinSketch: {}", - e - )) - })?; - Ok(Box::new(acc)) - } - - // Cardinality sketches - SketchType::HLL => { - let acc = HllAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize HLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - - // Sketches that aren't implemented yet - _ => Err(DataFusionError::NotImplemented(format!( - "Accumulator deserialization not implemented for: {:?}", - summary_type - ))), - } -} - -/// Serialize an accumulator to bytes (native format). -/// -/// This is a convenience wrapper that calls serialize_to_bytes on the accumulator. -pub fn serialize_accumulator(acc: &dyn AggregateCore) -> Vec { - acc.serialize_to_bytes() -} - -/// Serialize an accumulator to Arroyo-compatible bytes (MessagePack format). -/// -/// For accumulators whose native serialize_to_bytes already uses MessagePack, -/// this delegates to serialize_to_bytes. For others (SumAccumulator, -/// SetAggregatorAccumulator, MultipleIncreaseAccumulator), this uses -/// their serialize_to_bytes_arroyo method. -pub fn serialize_accumulator_arroyo(acc: &dyn AggregateCore) -> Vec { - // Try to downcast to types that have a separate arroyo format - if let Some(sum_acc) = acc.as_any().downcast_ref::() { - return sum_acc.serialize_to_bytes_arroyo(); - } - if let Some(set_acc) = acc.as_any().downcast_ref::() { - return set_acc.serialize_to_bytes_arroyo(); - } - if let Some(inc_acc) = acc.as_any().downcast_ref::() { - return inc_acc.serialize_to_bytes_arroyo(); - } - if let Some(ms_acc) = acc.as_any().downcast_ref::() { - return ms_acc.serialize_to_bytes_arroyo(); - } - // All other accumulators already use MessagePack in serialize_to_bytes - acc.serialize_to_bytes() -} - -/// Deserialize bytes to a SingleSubpopulationAggregate for querying. -/// -/// This function returns a trait object that supports the query method. -/// Only works for single-subpopulation accumulators (Sum, Increase, MinMax, etc.). -/// -/// Note: Uses Arroyo/MessagePack format. -pub fn deserialize_single_subpopulation( - bytes: &[u8], - summary_type: &SketchType, -) -> Result, DataFusionError> { - match summary_type { - SketchType::Sum => { - let acc = SumAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize Sum: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::Increase => Err(DataFusionError::NotImplemented( - "Increase Arroyo deserialization not implemented".to_string(), - )), - SketchType::MinMax => Err(DataFusionError::NotImplemented( - "MinMax Arroyo deserialization not implemented".to_string(), - )), - SketchType::KLL => { - let acc = - DatasketchesKLLAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize KLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::HLL => { - let acc = HllAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize HLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - _ => Err(DataFusionError::NotImplemented(format!( - "SingleSubpopulationAggregate deserialization not implemented for: {:?}", - summary_type - ))), - } -} - -/// Deserialize bytes to a MultipleSubpopulationAggregate for querying. -/// -/// This function returns a trait object that supports querying by sub-key. -/// Works for multi-population accumulators (Hydra types, CountMinSketch, etc.). -/// -/// Note: Uses Arroyo/MessagePack format. -pub fn deserialize_multiple_subpopulation( - bytes: &[u8], - summary_type: &SketchType, -) -> Result, DataFusionError> { - match summary_type { - SketchType::MultipleIncrease => { - let acc = - MultipleIncreaseAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize MultipleIncrease: {}", - e - )) - })?; - Ok(Box::new(acc)) - } - SketchType::MultipleSum => { - let acc = - MultipleSumAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize MultipleSum: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::HydraKLL => { - let acc = - HydraKllSketchAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize HydraKLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::CountMinSketch => { - let acc = - CountMinSketchAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize CountMinSketch: {}", - e - )) - })?; - Ok(Box::new(acc)) - } - _ => Err(DataFusionError::NotImplemented(format!( - "MultipleSubpopulationAggregate deserialization not implemented for: {:?}", - summary_type - ))), - } -} - -/// Deserialize bytes to a keys accumulator (DeltaSetAggregator/SetAggregator). -/// -/// Returns a boxed AggregateCore whose `get_keys()` method enumerates the sub-keys -/// stored in the accumulator. -pub fn deserialize_keys_accumulator( - bytes: &[u8], - summary_type: &SketchType, -) -> Result, DataFusionError> { - match summary_type { - SketchType::DeltaSetAggregator => { - let acc = DeltaSetAggregatorAccumulator::deserialize_from_bytes_arroyo(bytes).map_err( - |e| { - DataFusionError::Internal(format!( - "Failed to deserialize DeltaSetAggregator: {}", - e - )) - }, - )?; - Ok(Box::new(acc)) - } - SketchType::SetAggregator => { - let acc = - SetAggregatorAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize SetAggregator: {}", e)) - })?; - Ok(Box::new(acc)) - } - _ => Err(DataFusionError::NotImplemented(format!( - "Keys accumulator deserialization not supported for: {:?}", - summary_type - ))), - } -} - -#[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 { - rmp_serde::to_vec(&value).unwrap() - } - - #[test] - fn test_deserialize_sum_accumulator() { - // SumAccumulator::deserialize_from_bytes_arroyo expects MessagePack f64 - let bytes = serialize_f64_arroyo(42.0); - - let restored = deserialize_accumulator(&bytes, &SketchType::Sum).unwrap(); - - assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); - } - - #[test] - fn test_deserialize_minmax_returns_not_implemented() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_accumulator(&bytes, &SketchType::MinMax); - - assert!(result.is_err()); - } - - #[test] - fn test_deserialize_hll_round_trip() { - // HLL is now a supported path. Build a real accumulator, serialize it, - // then verify deserialize_accumulator and deserialize_single_subpopulation - // both reconstruct a working accumulator that reports the same estimate. - use crate::data_model::SerializableToSink; - use crate::precompute_operators::HllAccumulator; - let mut acc = HllAccumulator::new(14); - for i in 0..500 { - acc.update(i as f64); - } - let bytes = acc.serialize_to_bytes(); - let want = acc.estimate(); - - let core = deserialize_accumulator(&bytes, &SketchType::HLL).expect("HLL deserialize"); - assert_eq!(core.get_accumulator_type(), AggregationType::HLL); - let core_hll = core - .as_any() - .downcast_ref::() - .expect("downcast HllAccumulator"); - assert_eq!(core_hll.estimate(), want); - - let single = - deserialize_single_subpopulation(&bytes, &SketchType::HLL).expect("HLL single-pop"); - let via_query = single - .query( - promql_utilities::query_logics::enums::Statistic::Cardinality, - None, - ) - .expect("Cardinality query"); - assert_eq!(via_query, want); - } - - #[test] - fn test_deserialize_unsupported_type() { - // MinMax keys-accumulator path remains unimplemented; ensure the generic - // dispatch still errors out cleanly for sketch types we haven't wired. - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_accumulator(&bytes, &SketchType::MinMax); - - assert!(result.is_err()); - } -} diff --git a/asap-query-engine/src/engines/physical/conversion.rs b/asap-query-engine/src/engines/physical/conversion.rs deleted file mode 100644 index 26cde347..00000000 --- a/asap-query-engine/src/engines/physical/conversion.rs +++ /dev/null @@ -1,452 +0,0 @@ -//! Conversion Utilities Module -//! -//! Provides functions to convert between store results and Arrow RecordBatches. -//! This enables DataFusion physical operators to work with precomputed outputs. - -use arrow::array::{ArrayRef, BinaryBuilder, Float64Array, StringBuilder}; -use arrow::datatypes::{DataType, Field, Schema}; -use arrow::record_batch::RecordBatch; -use datafusion::error::DataFusionError; -use std::collections::HashMap; -use std::sync::Arc; - -use crate::data_model::KeyByLabelValues; -use crate::engines::physical::accumulator_serde::serialize_accumulator_arroyo; -use crate::stores::traits::TimestampedBucketsMap; - -/// Convert store query results to an Arrow RecordBatch. -/// -/// The output schema is: [label_columns..., sketch (Binary)] -/// Each row represents one group key with its serialized accumulator. -/// -/// # Arguments -/// * `store_result` - HashMap from group key to accumulators -/// * `label_names` - Names of the label columns (in order) -/// -/// # Returns -/// A RecordBatch with label columns (Utf8) and a sketch column (Binary) -pub fn store_result_to_record_batch( - store_result: &TimestampedBucketsMap, - label_names: &[String], -) -> Result { - // Build arrays for each label column - let mut label_builders: Vec = - label_names.iter().map(|_| StringBuilder::new()).collect(); - - // Build array for sketch column - let mut sketch_builder = BinaryBuilder::new(); - - for (key_opt, timestamped_buckets) in store_result { - // Handle each accumulator for this key - for (_timestamps, acc) in timestamped_buckets { - // Add label values - if let Some(key) = key_opt { - for (i, label_value) in key.labels.iter().enumerate() { - if i < label_builders.len() { - label_builders[i].append_value(label_value); - } - } - // Pad with empty strings if key has fewer labels than expected - for item in label_builders.iter_mut().skip(key.labels.len()) { - item.append_value(""); - } - } else { - // No key - use empty strings for all labels - for builder in &mut label_builders { - builder.append_value(""); - } - } - - // Serialize accumulator to Arroyo-compatible bytes (MessagePack) - // so downstream operators can deserialize with deserialize_from_bytes_arroyo - let bytes = serialize_accumulator_arroyo(acc.as_ref()); - sketch_builder.append_value(&bytes); - } - } - - // Build schema - let mut fields: Vec = label_names - .iter() - .map(|name| Field::new(name, DataType::Utf8, true)) - .collect(); - fields.push(Field::new("sketch", DataType::Binary, false)); - let schema = Arc::new(Schema::new(fields)); - - // Build columns - let mut columns: Vec = label_builders - .iter_mut() - .map(|b| Arc::new(b.finish()) as ArrayRef) - .collect(); - columns.push(Arc::new(sketch_builder.finish())); - - RecordBatch::try_new(schema, columns) - .map_err(|e| DataFusionError::Internal(format!("Failed to create RecordBatch: {}", e))) -} - -/// Convert a RecordBatch with inferred values back to a result map. -/// -/// The input schema is expected to be: [label_columns..., value_column (Float64)] -/// -/// # Arguments -/// * `batch` - RecordBatch with label columns and a value column -/// * `label_names` - Names of the label columns (to identify which columns are labels) -/// * `value_column` - Name of the column containing the inferred values -/// -/// # Returns -/// A HashMap from group key to the extracted value -pub fn record_batch_to_result_map( - batch: &RecordBatch, - label_names: &[&str], - value_column: &str, -) -> Result, f64>, DataFusionError> { - let mut result: HashMap, f64> = HashMap::new(); - - // Find the value column - let value_col_idx = batch - .schema() - .fields() - .iter() - .position(|f| f.name() == value_column) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "No '{}' column found in batch schema", - value_column - )) - })?; - - let value_array = batch - .column(value_col_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal(format!("'{}' column is not Float64", value_column)) - })?; - - // Find label column indices - let label_indices: Vec = label_names - .iter() - .filter_map(|name| { - batch - .schema() - .fields() - .iter() - .position(|f| f.name() == *name) - }) - .collect(); - - for row_idx in 0..batch.num_rows() { - // Extract label values for this row - let labels: Vec = label_indices - .iter() - .map(|&col_idx| { - let col = batch.column(col_idx); - // Try to extract string value - if let Some(str_array) = col.as_any().downcast_ref::() { - str_array.value(row_idx).to_string() - } else { - String::new() - } - }) - .collect(); - - let key = if labels.is_empty() || labels.iter().all(|l| l.is_empty()) { - None - } else { - Some(KeyByLabelValues { labels }) - }; - - let value = value_array.value(row_idx); - result.insert(key, value); - } - - Ok(result) -} - -/// Helper function to count total rows in store result -pub fn count_store_result_rows(store_result: &TimestampedBucketsMap) -> usize { - store_result.values().map(|v| v.len()).sum() -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::precompute_operators::SumAccumulator; - use crate::stores::traits::TimestampedBucket; - - fn make_bucket(acc: Arc) -> TimestampedBucket { - ((0, 0), acc) - } - - #[test] - fn test_store_result_to_record_batch_basic() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - - let key1 = KeyByLabelValues { - labels: vec!["host-a".to_string()], - }; - let acc1 = Arc::new(SumAccumulator::with_sum(100.0)) as Arc; - store_result.insert(Some(key1), vec![make_bucket(acc1)]); - - let key2 = KeyByLabelValues { - labels: vec!["host-b".to_string()], - }; - let acc2 = Arc::new(SumAccumulator::with_sum(200.0)) as Arc; - store_result.insert(Some(key2), vec![make_bucket(acc2)]); - - let label_names = vec!["host".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - - assert_eq!(batch.num_rows(), 2); - assert_eq!(batch.num_columns(), 2); // host, sketch - } - - #[test] - fn test_store_result_to_record_batch_multiple_labels() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - - let key1 = KeyByLabelValues { - labels: vec!["host-a".to_string(), "region-1".to_string()], - }; - let acc1 = Arc::new(SumAccumulator::with_sum(100.0)) as Arc; - store_result.insert(Some(key1), vec![make_bucket(acc1)]); - - let label_names = vec!["host".to_string(), "region".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - - assert_eq!(batch.num_rows(), 1); - assert_eq!(batch.num_columns(), 3); // host, region, sketch - } - - #[test] - fn test_store_result_to_record_batch_no_key() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - - let acc = Arc::new(SumAccumulator::with_sum(500.0)) as Arc; - store_result.insert(None, vec![make_bucket(acc)]); - - let label_names: Vec = vec![]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - - assert_eq!(batch.num_rows(), 1); - assert_eq!(batch.num_columns(), 1); // just sketch - } - - #[test] - fn test_record_batch_to_result_map() { - // Create a test batch with [host, value] - let host_array = arrow::array::StringArray::from(vec!["host-a", "host-b"]); - let value_array = Float64Array::from(vec![100.0, 200.0]); - - let schema = Arc::new(Schema::new(vec![ - Field::new("host", DataType::Utf8, true), - Field::new("value", DataType::Float64, false), - ])); - - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(host_array) as ArrayRef, - Arc::new(value_array) as ArrayRef, - ], - ) - .unwrap(); - - let result = record_batch_to_result_map(&batch, &["host"], "value").unwrap(); - - assert_eq!(result.len(), 2); - - let key_a = KeyByLabelValues { - labels: vec!["host-a".to_string()], - }; - assert_eq!(result.get(&Some(key_a)), Some(&100.0)); - - let key_b = KeyByLabelValues { - labels: vec!["host-b".to_string()], - }; - assert_eq!(result.get(&Some(key_b)), Some(&200.0)); - } - - // ======================================================================== - // Edge case tests - // ======================================================================== - - #[test] - fn test_store_result_to_record_batch_empty() { - let store_result: TimestampedBucketsMap = HashMap::new(); - let label_names = vec!["host".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - assert_eq!(batch.num_rows(), 0); - assert_eq!(batch.num_columns(), 2); // host + sketch - } - - #[test] - fn test_store_result_to_record_batch_five_labels() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - let key = KeyByLabelValues { - labels: vec![ - "a".to_string(), - "b".to_string(), - "c".to_string(), - "d".to_string(), - "e".to_string(), - ], - }; - store_result.insert( - Some(key), - vec![make_bucket( - Arc::new(SumAccumulator::with_sum(1.0)) as Arc - )], - ); - let label_names: Vec = vec!["l1", "l2", "l3", "l4", "l5"] - .into_iter() - .map(String::from) - .collect(); - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - assert_eq!(batch.num_rows(), 1); - assert_eq!(batch.num_columns(), 6); // 5 labels + sketch - } - - #[test] - fn test_store_result_to_record_batch_special_chars() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - let key = KeyByLabelValues { - labels: vec!["host,with,commas".to_string(), "région-1".to_string()], - }; - store_result.insert( - Some(key), - vec![make_bucket( - Arc::new(SumAccumulator::with_sum(42.0)) as Arc - )], - ); - let label_names = vec!["host".to_string(), "region".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - assert_eq!(batch.num_rows(), 1); - - // Verify the special characters survived - let host_col = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(host_col.value(0), "host,with,commas"); - let region_col = batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(region_col.value(0), "région-1"); - } - - #[test] - fn test_store_result_to_record_batch_multiple_timestamps_per_key() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - let key = KeyByLabelValues { - labels: vec!["host-a".to_string()], - }; - // 3 buckets for the same key - store_result.insert( - Some(key), - vec![ - ( - (100, 200), - Arc::new(SumAccumulator::with_sum(10.0)) as Arc, - ), - ( - (200, 300), - Arc::new(SumAccumulator::with_sum(20.0)) as Arc, - ), - ( - (300, 400), - Arc::new(SumAccumulator::with_sum(30.0)) as Arc, - ), - ], - ); - let label_names = vec!["host".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - assert_eq!(batch.num_rows(), 3, "3 buckets should produce 3 rows"); - } - - #[test] - fn test_record_batch_to_result_map_no_labels() { - let value_array = Float64Array::from(vec![42.0]); - let schema = Arc::new(Schema::new(vec![Field::new( - "value", - DataType::Float64, - false, - )])); - let batch = RecordBatch::try_new(schema, vec![Arc::new(value_array) as ArrayRef]).unwrap(); - - let result = record_batch_to_result_map(&batch, &[], "value").unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result.get(&None), Some(&42.0)); - } - - #[test] - fn test_record_batch_to_result_map_missing_value_column() { - let host_array = arrow::array::StringArray::from(vec!["host-a"]); - let schema = Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])); - let batch = RecordBatch::try_new(schema, vec![Arc::new(host_array) as ArrayRef]).unwrap(); - - let result = record_batch_to_result_map(&batch, &["host"], "value"); - assert!(result.is_err(), "Missing value column should produce error"); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains("value"), - "Error should mention 'value' column" - ); - } - - #[test] - fn test_record_batch_to_result_map_wrong_type() { - // value column is Utf8 instead of Float64 - let host_array = arrow::array::StringArray::from(vec!["host-a"]); - let value_array = arrow::array::StringArray::from(vec!["not_a_number"]); - let schema = Arc::new(Schema::new(vec![ - Field::new("host", DataType::Utf8, true), - Field::new("value", DataType::Utf8, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(host_array) as ArrayRef, - Arc::new(value_array) as ArrayRef, - ], - ) - .unwrap(); - - let result = record_batch_to_result_map(&batch, &["host"], "value"); - assert!(result.is_err(), "Wrong type for value column should error"); - } - - #[test] - fn test_count_store_result_rows() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - let key1 = KeyByLabelValues { - labels: vec!["a".to_string()], - }; - let key2 = KeyByLabelValues { - labels: vec!["b".to_string()], - }; - store_result.insert( - Some(key1), - vec![ - make_bucket( - Arc::new(SumAccumulator::with_sum(1.0)) as Arc - ), - make_bucket( - Arc::new(SumAccumulator::with_sum(2.0)) as Arc - ), - ], - ); - store_result.insert( - Some(key2), - vec![make_bucket( - Arc::new(SumAccumulator::with_sum(3.0)) as Arc - )], - ); - assert_eq!(count_store_result_rows(&store_result), 3); - - let empty: TimestampedBucketsMap = HashMap::new(); - assert_eq!(count_store_result_rows(&empty), 0); - } -} diff --git a/asap-query-engine/src/engines/physical/mod.rs b/asap-query-engine/src/engines/physical/mod.rs deleted file mode 100644 index 4cc5e54b..00000000 --- a/asap-query-engine/src/engines/physical/mod.rs +++ /dev/null @@ -1,29 +0,0 @@ -//! Physical Execution Operators for DataFusion -//! -//! This module provides physical execution plan operators that implement -//! DataFusion's ExecutionPlan trait for precomputed summary operations. - -pub mod accumulator_serde; -pub mod conversion; -pub mod planner; -pub mod precomputed_summary_read_exec; -pub mod summary_infer_exec; -pub mod summary_merge_multiple_exec; - -pub use planner::{CustomQueryPlanner, QueryEngineExtensionPlanner}; -pub use precomputed_summary_read_exec::PrecomputedSummaryReadExec; -pub use summary_infer_exec::SummaryInferExec; -pub use summary_merge_multiple_exec::SummaryMergeMultipleExec; - -use arrow::datatypes::SchemaRef; - -/// Format an Arrow schema as a compact string for debug logging. -/// Example: `{host: Utf8, region: Utf8, sketch: Binary}` -pub(crate) fn format_schema(schema: &SchemaRef) -> String { - let fields: Vec = schema - .fields() - .iter() - .map(|f| format!("{}: {:?}", f.name(), f.data_type())) - .collect(); - format!("{{{}}}", fields.join(", ")) -} diff --git a/asap-query-engine/src/engines/physical/planner.rs b/asap-query-engine/src/engines/physical/planner.rs deleted file mode 100644 index b0dc7b9e..00000000 --- a/asap-query-engine/src/engines/physical/planner.rs +++ /dev/null @@ -1,180 +0,0 @@ -//! Extension Planner for QueryEngineRust -//! -//! This module provides an ExtensionPlanner implementation that converts -//! custom logical operators (PrecomputedSummaryRead, SummaryMergeMultiple) -//! into their physical execution counterparts. - -use async_trait::async_trait; -use datafusion::error::DataFusionError; -use datafusion::execution::context::SessionState; -use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode}; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}; -use datafusion_summary_library::{ - PrecomputedSummaryRead, SketchType, SummaryInfer, SummaryMergeMultiple, -}; -use std::fmt; -use std::sync::Arc; - -use super::{PrecomputedSummaryReadExec, SummaryInferExec, SummaryMergeMultipleExec}; -use crate::stores::Store; - -/// Extension planner that handles custom logical operators for QueryEngineRust. -/// -/// This planner knows how to convert: -/// - PrecomputedSummaryRead -> PrecomputedSummaryReadExec -/// - SummaryMergeMultiple -> SummaryMergeMultipleExec -/// -/// Note: SummaryInfer is handled by datafusion_summary_library's planner -pub struct QueryEngineExtensionPlanner { - /// Reference to the store for reading precomputed outputs - store: Arc, -} - -impl QueryEngineExtensionPlanner { - pub fn new(store: Arc) -> Self { - Self { store } - } -} - -#[async_trait] -impl ExtensionPlanner for QueryEngineExtensionPlanner { - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - _logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - _session_state: &SessionState, - ) -> Result>, DataFusionError> { - // Try to downcast to PrecomputedSummaryRead - if let Some(read) = node.as_any().downcast_ref::() { - return Ok(Some(Arc::new(PrecomputedSummaryReadExec::new( - read.clone(), - self.store.clone(), - )))); - } - - // Try to downcast to SummaryMergeMultiple - if let Some(merge) = node.as_any().downcast_ref::() { - if physical_inputs.len() != 1 { - return Err(DataFusionError::Internal( - "SummaryMergeMultiple expects exactly one input".to_string(), - )); - } - return Ok(Some(Arc::new(SummaryMergeMultipleExec::new( - merge.clone(), - physical_inputs[0].clone(), - )))); - } - - // Try to downcast to SummaryInfer - if let Some(infer) = node.as_any().downcast_ref::() { - // Extract summary_type and sketch_column from the first logical input (values SummaryMergeMultiple) - let values_input_plan = _logical_inputs.first().ok_or_else(|| { - DataFusionError::Internal("SummaryInfer has no logical inputs".to_string()) - })?; - - let (summary_type, sketch_column) = extract_merge_info(values_input_plan, "values")?; - - if _logical_inputs.len() == 2 && physical_inputs.len() == 2 { - // Dual-input: extract keys summary type from second logical input - let keys_input_plan = _logical_inputs[1]; - let (keys_summary_type, _) = extract_merge_info(keys_input_plan, "keys")?; - - return Ok(Some(Arc::new(SummaryInferExec::new_dual_input( - infer.clone(), - physical_inputs[0].clone(), - physical_inputs[1].clone(), - summary_type, - keys_summary_type, - sketch_column, - )))); - } else if physical_inputs.len() == 1 { - // Single-input (original path) - return Ok(Some(Arc::new(SummaryInferExec::new( - infer.clone(), - physical_inputs[0].clone(), - summary_type, - sketch_column, - )))); - } else { - return Err(DataFusionError::Internal(format!( - "SummaryInfer: unexpected number of inputs: logical={}, physical={}", - _logical_inputs.len(), - physical_inputs.len() - ))); - } - } - - // Not a node we handle - let other planners try - Ok(None) - } -} - -/// Extract summary_type and sketch_column from a SummaryMergeMultiple logical plan node. -fn extract_merge_info( - plan: &LogicalPlan, - label: &str, -) -> Result<(SketchType, String), DataFusionError> { - match plan { - LogicalPlan::Extension(ext) => ext - .node - .as_any() - .downcast_ref::() - .map(|merge| { - ( - merge.summary_type().clone(), - merge.sketch_column().to_string(), - ) - }) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "SummaryInfer {} input is not SummaryMergeMultiple", - label - )) - }), - _ => Err(DataFusionError::Internal(format!( - "SummaryInfer {} input must be an Extension node", - label - ))), - } -} - -/// Custom query planner that combines the default DataFusion planner with -/// our QueryEngineExtensionPlanner for custom operators. -pub struct CustomQueryPlanner { - store: Arc, -} - -impl CustomQueryPlanner { - pub fn new(store: Arc) -> Self { - Self { store } - } -} - -impl fmt::Debug for CustomQueryPlanner { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CustomQueryPlanner") - .field("store", &"") - .finish() - } -} - -#[async_trait] -impl datafusion::execution::context::QueryPlanner for CustomQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &SessionState, - ) -> Result, DataFusionError> { - // Create default planner with our extension planner - let extension_planner = QueryEngineExtensionPlanner::new(self.store.clone()); - let physical_planner = - DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(extension_planner)]); - - physical_planner - .create_physical_plan(logical_plan, session_state) - .await - } -} diff --git a/asap-query-engine/src/engines/physical/precomputed_summary_read_exec.rs b/asap-query-engine/src/engines/physical/precomputed_summary_read_exec.rs deleted file mode 100644 index fa991ecd..00000000 --- a/asap-query-engine/src/engines/physical/precomputed_summary_read_exec.rs +++ /dev/null @@ -1,169 +0,0 @@ -//! PrecomputedSummaryReadExec - Physical execution operator for reading precomputed summaries -//! -//! This operator reads precomputed aggregates from a Store and produces -//! RecordBatches with label columns and a serialized sketch column. - -use arrow::datatypes::SchemaRef; -use datafusion::error::DataFusionError; -use datafusion::execution::TaskContext; -use datafusion::logical_expr::UserDefinedLogicalNodeCore; -use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion::physical_plan::{ - stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, - PlanProperties, SendableRecordBatchStream, -}; -use datafusion_summary_library::PrecomputedSummaryRead; -use futures::stream; -use std::any::Any; -use std::fmt; -use std::sync::Arc; -use std::time::Instant; -use tracing::debug; - -use super::format_schema; -use crate::engines::physical::conversion::store_result_to_record_batch; -use crate::stores::Store; - -/// Physical execution plan for reading precomputed summaries from a store. -pub struct PrecomputedSummaryReadExec { - /// The logical operator this was created from - logical_node: PrecomputedSummaryRead, - /// Reference to the store - store: Arc, - /// Output schema - schema: SchemaRef, - /// Plan properties (cached) - properties: PlanProperties, -} - -impl PrecomputedSummaryReadExec { - pub fn new(logical_node: PrecomputedSummaryRead, store: Arc) -> Self { - // Convert DFSchema to Schema - let schema = Arc::new(logical_node.schema().as_ref().into()); - - let properties = PlanProperties::new( - EquivalenceProperties::new(Arc::clone(&schema)), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - logical_node, - store, - schema, - properties, - } - } -} - -impl fmt::Debug for PrecomputedSummaryReadExec { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("PrecomputedSummaryReadExec") - .field("metric", &self.logical_node.metric()) - .field("aggregation_id", &self.logical_node.aggregation_id()) - .finish() - } -} - -impl DisplayAs for PrecomputedSummaryReadExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "PrecomputedSummaryReadExec: metric={}, agg_id={}, range=[{}, {}]", - self.logical_node.metric(), - self.logical_node.aggregation_id(), - self.logical_node.start_timestamp(), - self.logical_node.end_timestamp() - ) - } -} - -impl ExecutionPlan for PrecomputedSummaryReadExec { - fn name(&self) -> &str { - "PrecomputedSummaryReadExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } - - fn properties(&self) -> &PlanProperties { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![] // Leaf node - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result, DataFusionError> { - // No children to replace - Ok(self) - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - debug!( - metric = %self.logical_node.metric(), - aggregation_id = self.logical_node.aggregation_id(), - start_timestamp = self.logical_node.start_timestamp(), - end_timestamp = self.logical_node.end_timestamp(), - is_exact_query = self.logical_node.is_exact_query(), - output_schema = %format_schema(&self.schema), - output_labels = ?self.logical_node.output_labels(), - "PrecomputedSummaryReadExec::execute" - ); - - // Query the store - let store_query_start = Instant::now(); - let store_result = if self.logical_node.is_exact_query() { - self.store.query_precomputed_output_exact( - self.logical_node.metric(), - self.logical_node.aggregation_id(), - self.logical_node.start_timestamp(), - self.logical_node.end_timestamp(), - ) - } else { - self.store.query_precomputed_output( - self.logical_node.metric(), - self.logical_node.aggregation_id(), - self.logical_node.start_timestamp(), - self.logical_node.end_timestamp(), - ) - } - .map_err(DataFusionError::External)?; - debug!( - store_query_ms = format!("{:.2}", store_query_start.elapsed().as_secs_f64() * 1000.0), - unique_keys = store_result.len(), - "PrecomputedSummaryReadExec store query complete" - ); - - // Convert to RecordBatch - let convert_start = Instant::now(); - let label_names: Vec = self.logical_node.output_labels().to_vec(); - let batch = store_result_to_record_batch(&store_result, &label_names)?; - - debug!( - convert_ms = format!("{:.2}", convert_start.elapsed().as_secs_f64() * 1000.0), - output_rows = batch.num_rows(), - output_cols = batch.num_columns(), - "PrecomputedSummaryReadExec produced batch" - ); - - // Create a stream that yields this single batch - let schema = self.schema.clone(); - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema, - stream::once(async move { Ok(batch) }), - ))) - } -} diff --git a/asap-query-engine/src/engines/physical/summary_infer_exec.rs b/asap-query-engine/src/engines/physical/summary_infer_exec.rs deleted file mode 100644 index 52885528..00000000 --- a/asap-query-engine/src/engines/physical/summary_infer_exec.rs +++ /dev/null @@ -1,769 +0,0 @@ -//! SummaryInferExec - Physical execution operator for extracting values from summaries -//! -//! This operator extracts values from serialized accumulators (summaries). -//! -//! **Single-population path** (no keys_input): -//! Input: rows with [label columns, sketch column] -//! Output: rows with [label columns, value column] -//! -//! **Multi-population path** (keys_input present): -//! Input 0 (values): rows with [spatial_label columns, sketch column] -//! Input 1 (keys): rows with [spatial_label columns, sketch column] -//! Output: rows with [spatial_label columns, sub_key columns, value column] - -use arrow::array::{ArrayRef, BinaryArray, Float64Builder, StringBuilder}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use arrow::record_batch::RecordBatch; -use datafusion::error::DataFusionError; -use datafusion::execution::TaskContext; -use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion::physical_plan::common::collect; -use datafusion::physical_plan::{ - stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, - PlanProperties, SendableRecordBatchStream, -}; -use datafusion_summary_library::{InferOperation, SketchType, SummaryInfer}; -use futures::stream; -use promql_utilities::query_logics::enums::Statistic; -use std::any::Any; -use std::collections::HashMap; -use std::fmt; -use std::sync::Arc; -use std::time::Instant; -use tracing::debug; - -use super::format_schema; -use crate::engines::physical::accumulator_serde::{ - deserialize_accumulator, deserialize_keys_accumulator, deserialize_multiple_subpopulation, - deserialize_single_subpopulation, -}; - -/// Physical execution plan for extracting values from serialized summaries. -pub struct SummaryInferExec { - /// The logical operator this was created from - logical_node: SummaryInfer, - /// Input execution plan (values branch) - input: Arc, - /// Optional second input (keys branch) for multi-population accumulators - keys_input: Option>, - /// Type of summary being inferred (determines deserialization) - summary_type: SketchType, - /// Type of the keys summary (only set for dual-input) - keys_summary_type: Option, - /// Name of the sketch column in the input schema - sketch_column: String, - /// Output schema (labels + value) - schema: SchemaRef, - /// Plan properties (cached) - properties: PlanProperties, -} - -impl SummaryInferExec { - pub fn new( - logical_node: SummaryInfer, - input: Arc, - summary_type: SketchType, - sketch_column: String, - ) -> Self { - let schema = Self::build_schema(&logical_node, &input, None, &sketch_column); - - let properties = PlanProperties::new( - EquivalenceProperties::new(Arc::clone(&schema)), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - logical_node, - input, - keys_input: None, - summary_type, - keys_summary_type: None, - sketch_column, - schema, - properties, - } - } - - /// Create a dual-input SummaryInferExec for multi-population accumulators. - pub fn new_dual_input( - logical_node: SummaryInfer, - input: Arc, - keys_input: Arc, - summary_type: SketchType, - keys_summary_type: SketchType, - sketch_column: String, - ) -> Self { - let schema = Self::build_schema(&logical_node, &input, Some(&keys_input), &sketch_column); - - let properties = PlanProperties::new( - EquivalenceProperties::new(Arc::clone(&schema)), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - logical_node, - input, - keys_input: Some(keys_input), - summary_type, - keys_summary_type: Some(keys_summary_type), - sketch_column, - schema, - properties, - } - } - - fn build_schema( - logical_node: &SummaryInfer, - input: &Arc, - _keys_input: Option<&Arc>, - sketch_column: &str, - ) -> SchemaRef { - let input_schema = input.schema(); - - // Build output schema: label columns from input (minus sketch) ... - let mut fields: Vec = input_schema - .fields() - .iter() - .filter(|f| f.name() != sketch_column) - .map(|f| f.as_ref().clone()) - .collect(); - - // ... plus sub-key columns for dual-input (group_key_columns) - for key_col in &logical_node.group_key_columns { - fields.push(Field::new(key_col, DataType::Utf8, true)); - } - - // ... plus output value columns (one per operation) - for output_name in &logical_node.output_names { - fields.push(Field::new(output_name, DataType::Float64, false)); - } - - Arc::new(Schema::new(fields)) - } - - /// Map InferOperation to Statistic for accumulator query - fn infer_op_to_statistic(op: &InferOperation) -> Statistic { - match op { - InferOperation::ExtractSum => Statistic::Sum, - InferOperation::ExtractCount => Statistic::Count, - InferOperation::ExtractMin => Statistic::Min, - InferOperation::ExtractMax => Statistic::Max, - InferOperation::ExtractIncrease => Statistic::Increase, - InferOperation::ExtractRate => Statistic::Rate, - InferOperation::CountDistinct => Statistic::Cardinality, - InferOperation::Quantile(_) | InferOperation::Median => Statistic::Quantile, - InferOperation::TopK(_) => Statistic::Topk, - // All other operations - use Count as fallback - _ => Statistic::Count, - } - } - - /// Extract query kwargs from an InferOperation. - /// - /// Some operations embed parameters (e.g. Quantile embeds the quantile value, - /// TopK embeds k) that accumulators need via the kwargs HashMap. - fn infer_op_to_kwargs(op: &InferOperation) -> Option> { - match op { - InferOperation::Quantile(q_u16) => { - let q = *q_u16 as f64 / 10000.0; - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), q.to_string()); - Some(kwargs) - } - InferOperation::Median => { - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.5".to_string()); - Some(kwargs) - } - InferOperation::TopK(k) => { - let mut kwargs = HashMap::new(); - kwargs.insert("k".to_string(), k.to_string()); - Some(kwargs) - } - _ => None, - } - } -} - -impl fmt::Debug for SummaryInferExec { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("SummaryInferExec") - .field("operations", &self.logical_node.operations) - .field("has_keys_input", &self.keys_input.is_some()) - .finish() - } -} - -impl DisplayAs for SummaryInferExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "SummaryInferExec: operations={:?}, dual_input={}", - self.logical_node.operations, - self.keys_input.is_some() - ) - } -} - -impl ExecutionPlan for SummaryInferExec { - fn name(&self) -> &str { - "SummaryInferExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } - - fn properties(&self) -> &PlanProperties { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - let mut children = vec![&self.input]; - if let Some(ref keys_input) = self.keys_input { - children.push(keys_input); - } - children - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result, DataFusionError> { - match (children.len(), &self.keys_input) { - (1, None) => Ok(Arc::new(Self::new( - self.logical_node.clone(), - children[0].clone(), - self.summary_type.clone(), - self.sketch_column.clone(), - ))), - (2, Some(_)) => Ok(Arc::new(Self::new_dual_input( - self.logical_node.clone(), - children[0].clone(), - children[1].clone(), - self.summary_type.clone(), - self.keys_summary_type.clone().unwrap(), - self.sketch_column.clone(), - ))), - _ => Err(DataFusionError::Internal(format!( - "SummaryInferExec: expected {} children, got {}", - if self.keys_input.is_some() { 2 } else { 1 }, - children.len() - ))), - } - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let schema = self.schema.clone(); - let schema_for_stream = schema.clone(); - let logical_node = self.logical_node.clone(); - let summary_type = self.summary_type.clone(); - let sketch_column = self.sketch_column.clone(); - - debug!( - input_schema = %format_schema(&self.input.schema()), - output_schema = %format_schema(&self.schema), - operations = ?self.logical_node.operations, - sketch_column = %self.sketch_column, - has_keys_input = self.keys_input.is_some(), - "SummaryInferExec::execute" - ); - - if let Some(ref keys_input) = self.keys_input { - // Multi-population path - let values_stream = self.input.execute(partition, context.clone())?; - let keys_stream = keys_input.execute(partition, context)?; - let keys_summary_type = self.keys_summary_type.clone().unwrap(); - - let output_stream = async move { - let collect_start = Instant::now(); - let values_batches = collect(values_stream).await?; - let keys_batches = collect(keys_stream).await?; - let values_rows: usize = values_batches.iter().map(|b| b.num_rows()).sum(); - let keys_rows: usize = keys_batches.iter().map(|b| b.num_rows()).sum(); - debug!( - collect_ms = format!("{:.2}", collect_start.elapsed().as_secs_f64() * 1000.0), - values_rows, keys_rows, "SummaryInferExec collected dual-input batches" - ); - - let infer_start = Instant::now(); - let result = process_dual_input( - &values_batches, - &keys_batches, - &logical_node, - &summary_type, - &keys_summary_type, - &sketch_column, - &schema, - )?; - debug!( - infer_ms = format!("{:.2}", infer_start.elapsed().as_secs_f64() * 1000.0), - output_rows = result.num_rows(), - "SummaryInferExec dual-input infer complete" - ); - Ok(result) - }; - - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema_for_stream, - stream::once(output_stream), - ))) - } else { - // Single-population path (original behavior) - let input_stream = self.input.execute(partition, context)?; - - let output_stream = async move { - let collect_start = Instant::now(); - let batches = collect(input_stream).await?; - let total_input_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - debug!( - collect_ms = format!("{:.2}", collect_start.elapsed().as_secs_f64() * 1000.0), - total_input_rows, - num_batches = batches.len(), - "SummaryInferExec collected input (single-pop)" - ); - - let mut all_label_values: Vec> = Vec::new(); - let mut all_result_values: Vec = Vec::new(); - - let self_keyed = is_self_keyed_multi_pop(&summary_type); - - let infer_start = Instant::now(); - for batch in &batches { - if self_keyed { - process_self_keyed_multi_pop_batch( - &mut all_label_values, - &mut all_result_values, - batch, - &logical_node, - &summary_type, - &sketch_column, - )?; - } else { - process_single_pop_batch( - &mut all_label_values, - &mut all_result_values, - batch, - &logical_node, - &summary_type, - &sketch_column, - )?; - } - } - debug!( - infer_ms = format!("{:.2}", infer_start.elapsed().as_secs_f64() * 1000.0), - output_rows = all_result_values.len(), - self_keyed, - "SummaryInferExec infer complete (single-pop)" - ); - - build_output_batch(&all_label_values, &all_result_values, &schema) - }; - - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema_for_stream, - stream::once(output_stream), - ))) - } - } -} - -// ============================================================================ -// Single-population processing (unchanged logic) -// ============================================================================ - -/// Process a batch and extract values from single-population accumulators -fn process_single_pop_batch( - all_label_values: &mut Vec>, - all_result_values: &mut Vec, - batch: &RecordBatch, - logical_node: &SummaryInfer, - summary_type: &SketchType, - sketch_column: &str, -) -> Result<(), DataFusionError> { - let sketch_idx = find_sketch_column_index(batch, sketch_column)?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Internal("Sketch column is not Binary".to_string()))?; - - let label_indices: Vec = (0..batch.num_columns()) - .filter(|&i| i != sketch_idx) - .collect(); - - let operation = logical_node - .operations - .first() - .ok_or_else(|| DataFusionError::Internal("SummaryInfer has no operations".to_string()))?; - - let statistic = SummaryInferExec::infer_op_to_statistic(operation); - let query_kwargs = SummaryInferExec::infer_op_to_kwargs(operation); - - for row in 0..batch.num_rows() { - let label_values = extract_label_values(batch, &label_indices, row); - let sketch_bytes = sketch_array.value(row); - - let accumulator = - deserialize_single_subpopulation(sketch_bytes, summary_type).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize accumulator: {}", e)) - })?; - - let value = accumulator - .query(statistic, query_kwargs.as_ref()) - .map_err(|e| { - DataFusionError::Internal(format!("Failed to query accumulator: {}", e)) - })?; - - all_label_values.push(label_values); - all_result_values.push(value); - } - - Ok(()) -} - -// ============================================================================ -// Self-keyed multi-population processing (single-input, keys from accumulator) -// ============================================================================ - -/// Returns true if the SketchType is a multi-population accumulator that -/// carries its own keys (via `get_keys()`), so it can be processed in -/// single-input mode without a separate keys stream. -fn is_self_keyed_multi_pop(summary_type: &SketchType) -> bool { - matches!( - summary_type, - SketchType::MultipleIncrease | SketchType::MultipleSum | SketchType::MultipleMinMax - ) -} - -/// Process a batch of self-keyed multi-population accumulators. -/// -/// For each row (spatial group): -/// 1. Deserialize as AggregateCore to call get_keys() -/// 2. Deserialize as MultipleSubpopulationAggregate to call query(stat, key) -/// 3. Emit one output row per sub-key -fn process_self_keyed_multi_pop_batch( - all_label_values: &mut Vec>, - all_result_values: &mut Vec, - batch: &RecordBatch, - logical_node: &SummaryInfer, - summary_type: &SketchType, - sketch_column: &str, -) -> Result<(), DataFusionError> { - let sketch_idx = find_sketch_column_index(batch, sketch_column)?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Internal("Sketch column is not Binary".to_string()))?; - - let label_indices: Vec = (0..batch.num_columns()) - .filter(|&i| i != sketch_idx) - .collect(); - - let operation = logical_node - .operations - .first() - .ok_or_else(|| DataFusionError::Internal("SummaryInfer has no operations".to_string()))?; - - let statistic = SummaryInferExec::infer_op_to_statistic(operation); - let query_kwargs = SummaryInferExec::infer_op_to_kwargs(operation); - - let num_sub_key_cols = logical_node.group_key_columns.len(); - - for row in 0..batch.num_rows() { - let spatial_labels = extract_label_values(batch, &label_indices, row); - let sketch_bytes = sketch_array.value(row); - - // Get keys from the accumulator itself - let acc = deserialize_accumulator(sketch_bytes, summary_type).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize accumulator: {}", e)) - })?; - let sub_keys = acc.get_keys().unwrap_or_default(); - - // Deserialize as multi-pop for querying - let multi_acc = - deserialize_multiple_subpopulation(sketch_bytes, summary_type).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize multi-pop accumulator: {}", - e - )) - })?; - - for sub_key in &sub_keys { - let value = multi_acc - .query(statistic, sub_key, query_kwargs.as_ref()) - .map_err(|e| { - DataFusionError::Internal(format!( - "Failed to query multi-pop accumulator for key {:?}: {}", - sub_key, e - )) - })?; - - let mut row_labels = spatial_labels.clone(); - for i in 0..num_sub_key_cols { - if i < sub_key.labels.len() { - row_labels.push(sub_key.labels[i].clone()); - } else { - row_labels.push(String::new()); - } - } - - all_label_values.push(row_labels); - all_result_values.push(value); - } - } - - Ok(()) -} - -// ============================================================================ -// Multi-population (dual-input) processing -// ============================================================================ - -/// Process dual-input: values + keys batches for multi-population accumulators. -/// -/// For each spatial group (row) in the values stream: -/// 1. Deserialize the value sketch as MultipleSubpopulationAggregate -/// 2. Find the matching spatial group in the keys stream -/// 3. Deserialize the keys accumulator, call get_keys() to enumerate sub-keys -/// 4. For each sub-key, call value_acc.query(statistic, sub_key) -> one output row -fn process_dual_input( - values_batches: &[RecordBatch], - keys_batches: &[RecordBatch], - logical_node: &SummaryInfer, - values_summary_type: &SketchType, - keys_summary_type: &SketchType, - sketch_column: &str, - schema: &SchemaRef, -) -> Result { - let operation = logical_node - .operations - .first() - .ok_or_else(|| DataFusionError::Internal("SummaryInfer has no operations".to_string()))?; - let statistic = SummaryInferExec::infer_op_to_statistic(operation); - let query_kwargs = SummaryInferExec::infer_op_to_kwargs(operation); - - // Build a lookup from spatial labels -> keys sketch bytes - // Key: spatial label values (as Vec), Value: serialized keys accumulator bytes - let keys_lookup = build_keys_lookup(keys_batches, sketch_column)?; - - let num_sub_key_cols = logical_node.group_key_columns.len(); - - let mut all_label_values: Vec> = Vec::new(); - let mut all_result_values: Vec = Vec::new(); - - for batch in values_batches { - let sketch_idx = find_sketch_column_index(batch, sketch_column)?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal("Values sketch column is not Binary".to_string()) - })?; - - let label_indices: Vec = (0..batch.num_columns()) - .filter(|&i| i != sketch_idx) - .collect(); - - for row in 0..batch.num_rows() { - let spatial_labels = extract_label_values(batch, &label_indices, row); - let value_sketch_bytes = sketch_array.value(row); - - // Deserialize the value sketch as MultipleSubpopulationAggregate - let value_acc = - deserialize_multiple_subpopulation(value_sketch_bytes, values_summary_type) - .map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize multi-pop value accumulator: {}", - e - )) - })?; - - // Find matching keys accumulator - let keys_bytes = keys_lookup.get(&spatial_labels).ok_or_else(|| { - DataFusionError::Internal(format!( - "No keys accumulator found for spatial group: {:?}", - spatial_labels - )) - })?; - - let keys_acc = - deserialize_keys_accumulator(keys_bytes, keys_summary_type).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize keys accumulator: {}", - e - )) - })?; - - let sub_keys = keys_acc.get_keys().unwrap_or_default(); - - debug!( - spatial_labels = ?spatial_labels, - num_sub_keys = sub_keys.len(), - "Processing multi-pop spatial group" - ); - - // For each sub-key, query the value accumulator - for sub_key in &sub_keys { - let value = value_acc - .query(statistic, sub_key, query_kwargs.as_ref()) - .map_err(|e| { - DataFusionError::Internal(format!( - "Failed to query multi-pop accumulator for key {:?}: {}", - sub_key, e - )) - })?; - - // Output row: [spatial_labels..., sub_key_labels..., value] - let mut row_labels = spatial_labels.clone(); - // Append sub-key label values - // sub_key.labels is Vec of values - for i in 0..num_sub_key_cols { - if i < sub_key.labels.len() { - row_labels.push(sub_key.labels[i].clone()); - } else { - row_labels.push(String::new()); - } - } - - all_label_values.push(row_labels); - all_result_values.push(value); - } - } - } - - debug!( - output_rows = all_result_values.len(), - "SummaryInferExec building output (multi-pop)" - ); - - build_output_batch(&all_label_values, &all_result_values, schema) -} - -/// Build a lookup map from spatial label values to keys sketch bytes. -fn build_keys_lookup( - keys_batches: &[RecordBatch], - sketch_column: &str, -) -> Result, Vec>, DataFusionError> { - let mut lookup: HashMap, Vec> = HashMap::new(); - - for batch in keys_batches { - let sketch_idx = find_sketch_column_index(batch, sketch_column)?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal("Keys sketch column is not Binary".to_string()) - })?; - - let label_indices: Vec = (0..batch.num_columns()) - .filter(|&i| i != sketch_idx) - .collect(); - - for row in 0..batch.num_rows() { - let label_values = extract_label_values(batch, &label_indices, row); - let sketch_bytes = sketch_array.value(row).to_vec(); - lookup.insert(label_values, sketch_bytes); - } - } - - Ok(lookup) -} - -// ============================================================================ -// Common helpers -// ============================================================================ - -/// Find the index of the sketch column in a batch -fn find_sketch_column_index( - batch: &RecordBatch, - sketch_column: &str, -) -> Result { - batch - .schema() - .fields() - .iter() - .position(|f| f.name() == sketch_column) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "Sketch column '{}' not found in batch schema: {:?}", - sketch_column, - batch - .schema() - .fields() - .iter() - .map(|f| f.name()) - .collect::>() - )) - }) -} - -/// Extract label values from a batch row -fn extract_label_values(batch: &RecordBatch, label_indices: &[usize], row: usize) -> Vec { - label_indices - .iter() - .map(|&idx| { - let col = batch.column(idx); - if let Some(str_array) = col.as_any().downcast_ref::() { - str_array.value(row).to_string() - } else { - String::new() - } - }) - .collect() -} - -/// Build output batch from extracted values -fn build_output_batch( - all_label_values: &[Vec], - all_result_values: &[f64], - schema: &SchemaRef, -) -> Result { - // Get number of label columns (schema fields minus the value column) - let num_label_cols = schema.fields().len() - 1; - - // Build label column builders - let mut label_builders: Vec = - (0..num_label_cols).map(|_| StringBuilder::new()).collect(); - - // Build value column - let mut value_builder = Float64Builder::new(); - - for (label_values, value) in all_label_values.iter().zip(all_result_values.iter()) { - // Add label values - for (i, label_value) in label_values.iter().enumerate() { - if i < label_builders.len() { - label_builders[i].append_value(label_value); - } - } - // Add value - value_builder.append_value(*value); - } - - // Build columns - let mut columns: Vec = label_builders - .iter_mut() - .map(|b| Arc::new(b.finish()) as ArrayRef) - .collect(); - columns.push(Arc::new(value_builder.finish())); - - RecordBatch::try_new(schema.clone(), columns) - .map_err(|e| DataFusionError::Internal(format!("Failed to build output batch: {}", e))) -} 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 deleted file mode 100644 index de5ec43d..00000000 --- a/asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs +++ /dev/null @@ -1,556 +0,0 @@ -//! SummaryMergeMultipleExec - Physical execution operator for merging summaries -//! -//! This operator merges multiple summaries with the same group key into one. -//! Input: multiple rows per group key with serialized accumulators -//! Output: one row per group key with merged accumulator - -use arrow::array::{ArrayRef, BinaryArray, BinaryBuilder, StringBuilder}; -use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; -use datafusion::error::DataFusionError; -use datafusion::execution::TaskContext; -use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion::physical_plan::common::collect; -use datafusion::physical_plan::{ - stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, - PlanProperties, SendableRecordBatchStream, -}; -use datafusion_summary_library::SummaryMergeMultiple; -use futures::stream; -use std::any::Any; -use std::collections::HashMap; -use std::fmt; -use std::sync::Arc; -use std::time::Instant; -use tracing::debug; - -use super::format_schema; -use crate::engines::physical::accumulator_serde::{ - deserialize_accumulator, serialize_accumulator_arroyo, -}; - -/// Physical execution plan for merging multiple summaries by group key. -pub struct SummaryMergeMultipleExec { - /// The logical operator this was created from - logical_node: SummaryMergeMultiple, - /// Input execution plan - input: Arc, - /// Output schema (same as input) - schema: SchemaRef, - /// Plan properties (cached) - properties: PlanProperties, -} - -impl SummaryMergeMultipleExec { - pub fn new(logical_node: SummaryMergeMultiple, input: Arc) -> Self { - let schema = input.schema(); - - let properties = PlanProperties::new( - EquivalenceProperties::new(Arc::clone(&schema)), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - logical_node, - input, - schema, - properties, - } - } -} - -impl fmt::Debug for SummaryMergeMultipleExec { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("SummaryMergeMultipleExec") - .field("group_by", &self.logical_node.group_by()) - .field("summary_type", &self.logical_node.summary_type()) - .finish() - } -} - -impl DisplayAs for SummaryMergeMultipleExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "SummaryMergeMultipleExec: group_by=[{}], type={}", - self.logical_node.group_by().join(", "), - self.logical_node.summary_type() - ) - } -} - -impl ExecutionPlan for SummaryMergeMultipleExec { - fn name(&self) -> &str { - "SummaryMergeMultipleExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } - - fn properties(&self) -> &PlanProperties { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result, DataFusionError> { - if children.len() != 1 { - return Err(DataFusionError::Internal( - "SummaryMergeMultipleExec expects exactly one child".to_string(), - )); - } - Ok(Arc::new(Self::new( - self.logical_node.clone(), - children[0].clone(), - ))) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let input_stream = self.input.execute(partition, context)?; - let schema = self.schema.clone(); - let schema_for_stream = schema.clone(); - let logical_node = self.logical_node.clone(); - - debug!( - input_schema = %format_schema(&self.input.schema()), - output_schema = %format_schema(&self.schema), - group_by = ?self.logical_node.group_by(), - sketch_column = %self.logical_node.sketch_column(), - summary_type = ?self.logical_node.summary_type(), - "SummaryMergeMultipleExec::execute" - ); - - // Use an async block to process all batches and merge - let output_stream = async move { - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - - // Collect all batches from input using datafusion's collect helper - let collect_start = Instant::now(); - let batches = collect(input_stream).await?; - let total_input_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - debug!( - collect_ms = format!("{:.2}", collect_start.elapsed().as_secs_f64() * 1000.0), - total_input_rows, - num_batches = batches.len(), - "SummaryMergeMultipleExec collected input" - ); - - // Process each batch - let merge_start = Instant::now(); - for batch in &batches { - process_batch(&mut groups, batch, &logical_node)?; - } - debug!( - merge_ms = format!("{:.2}", merge_start.elapsed().as_secs_f64() * 1000.0), - output_groups = groups.len(), - "SummaryMergeMultipleExec merged into groups" - ); - - // Build output batch - build_output_batch(&groups, &logical_node, &schema) - }; - - // Convert to stream - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema_for_stream, - stream::once(output_stream), - ))) - } -} - -/// Process a batch and accumulate into groups -fn process_batch( - groups: &mut HashMap, (Vec, Vec)>, - batch: &RecordBatch, - logical_node: &SummaryMergeMultiple, -) -> Result<(), DataFusionError> { - let group_by_cols = logical_node.group_by(); - let sketch_col_name = logical_node.sketch_column(); - - // Find column indices - let group_indices: Vec = group_by_cols - .iter() - .filter_map(|name| { - batch - .schema() - .fields() - .iter() - .position(|f| f.name() == name) - }) - .collect(); - - let sketch_idx = batch - .schema() - .fields() - .iter() - .position(|f| f.name() == sketch_col_name) - .ok_or_else(|| { - DataFusionError::Internal(format!("Sketch column '{}' not found", sketch_col_name)) - })?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Internal("Sketch column is not Binary".to_string()))?; - - for row in 0..batch.num_rows() { - // Extract group key - let group_key: Vec = group_indices - .iter() - .map(|&idx| { - let col = batch.column(idx); - if let Some(str_array) = col.as_any().downcast_ref::() { - str_array.value(row).to_string() - } else { - String::new() - } - }) - .collect(); - - // Get sketch bytes - let sketch_bytes = sketch_array.value(row); - - // Merge with existing group or insert new - if let Some((_, existing_bytes)) = groups.get_mut(&group_key) { - // Deserialize both accumulators and merge - let existing_acc = - deserialize_accumulator(existing_bytes, logical_node.summary_type())?; - let new_acc = deserialize_accumulator(sketch_bytes, logical_node.summary_type())?; - - // Merge accumulators - let merged = existing_acc.merge_with(new_acc.as_ref()).map_err(|e| { - DataFusionError::Internal(format!("Failed to merge accumulators: {}", e)) - })?; - - // Serialize merged accumulator in arroyo format for downstream deserialization - *existing_bytes = serialize_accumulator_arroyo(merged.as_ref()); - } else { - // First time seeing this group - groups.insert(group_key.clone(), (group_key, sketch_bytes.to_vec())); - } - } - - Ok(()) -} - -/// Build output batch from merged groups (public for testing) -pub(crate) fn build_output_batch( - groups: &HashMap, (Vec, Vec)>, - logical_node: &SummaryMergeMultiple, - schema: &SchemaRef, -) -> Result { - let group_by_cols = logical_node.group_by(); - - // Build arrays for each column - let mut label_builders: Vec = - group_by_cols.iter().map(|_| StringBuilder::new()).collect(); - let mut sketch_builder = BinaryBuilder::new(); - - for (label_values, bytes) in groups.values() { - // Add label values - for (i, value) in label_values.iter().enumerate() { - if i < label_builders.len() { - label_builders[i].append_value(value); - } - } - // Add sketch bytes - sketch_builder.append_value(bytes); - } - - // Build columns - let mut columns: Vec = label_builders - .iter_mut() - .map(|b| Arc::new(b.finish()) as ArrayRef) - .collect(); - columns.push(Arc::new(sketch_builder.finish())); - - RecordBatch::try_new(schema.clone(), columns) - .map_err(|e| DataFusionError::Internal(format!("Failed to build output batch: {}", e))) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::data_model::{AggregationType, KeyByLabelValues}; - use crate::engines::physical::accumulator_serde::serialize_accumulator_arroyo; - use crate::precompute_operators::{ - DatasketchesKLLAccumulator, SetAggregatorAccumulator, SumAccumulator, - }; - use arrow::array::StringArray; - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_summary_library::SketchType; - - /// Helper to create a RecordBatch with [host (Utf8), sketch (Binary)] - fn make_batch(rows: Vec<(&str, Vec)>) -> RecordBatch { - let mut host_builder = StringBuilder::new(); - let mut sketch_builder = BinaryBuilder::new(); - for (host, sketch_bytes) in &rows { - host_builder.append_value(host); - sketch_builder.append_value(sketch_bytes); - } - let schema = Arc::new(Schema::new(vec![ - Field::new("host", DataType::Utf8, true), - Field::new("sketch", DataType::Binary, false), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(host_builder.finish()) as ArrayRef, - Arc::new(sketch_builder.finish()) as ArrayRef, - ], - ) - .unwrap() - } - - /// Helper to create a SummaryMergeMultiple logical node for testing - fn make_logical_node(summary_type: SketchType) -> SummaryMergeMultiple { - use arrow::datatypes::DataType as DT; - use datafusion::common::DFSchema; - use datafusion::logical_expr::{Extension, LogicalPlan}; - use datafusion_summary_library::PrecomputedSummaryRead; - - let fields = vec![ - (None, Arc::new(Field::new("host", DT::Utf8, true))), - (None, Arc::new(Field::new("sketch", DT::Binary, false))), - ]; - let schema = Arc::new(DFSchema::new_with_metadata(fields, Default::default()).unwrap()); - let read = PrecomputedSummaryRead::new( - "test".to_string(), - 1, - 0, - 1000, - true, - vec!["host".to_string()], - summary_type.clone(), - schema, - ); - let read_plan = LogicalPlan::Extension(Extension { - node: Arc::new(read), - }); - SummaryMergeMultiple::new( - Arc::new(read_plan), - vec!["host".to_string()], - "sketch".to_string(), - summary_type, - ) - } - - #[test] - fn test_merge_single_row_passthrough() { - let acc = SumAccumulator::with_sum(42.0); - let bytes = serialize_accumulator_arroyo(&acc); - let batch = make_batch(vec![("host-a", bytes.clone())]); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - // 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(), AggregationType::Sum); - } - - #[test] - fn test_merge_two_sums_same_group() { - let acc1 = SumAccumulator::with_sum(50.0); - let acc2 = SumAccumulator::with_sum(50.0); - let bytes1 = serialize_accumulator_arroyo(&acc1); - let bytes2 = serialize_accumulator_arroyo(&acc2); - let batch = make_batch(vec![("host-a", bytes1), ("host-a", bytes2)]); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - let (_, merged_bytes) = groups.values().next().unwrap(); - let restored = - crate::engines::physical::accumulator_serde::deserialize_single_subpopulation( - merged_bytes, - &SketchType::Sum, - ) - .unwrap(); - let value = restored - .query(promql_utilities::query_logics::enums::Statistic::Sum, None) - .unwrap(); - assert!( - (value - 100.0).abs() < 1e-10, - "Merged sum should be 100.0, got {}", - value - ); - } - - #[test] - fn test_merge_three_sums_associativity() { - let bytes: Vec> = [30.0, 40.0, 30.0] - .iter() - .map(|v| serialize_accumulator_arroyo(&SumAccumulator::with_sum(*v))) - .collect(); - let batch = make_batch(vec![ - ("host-a", bytes[0].clone()), - ("host-a", bytes[1].clone()), - ("host-a", bytes[2].clone()), - ]); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - let (_, merged_bytes) = groups.values().next().unwrap(); - let restored = - crate::engines::physical::accumulator_serde::deserialize_single_subpopulation( - merged_bytes, - &SketchType::Sum, - ) - .unwrap(); - let value = restored - .query(promql_utilities::query_logics::enums::Statistic::Sum, None) - .unwrap(); - assert!( - (value - 100.0).abs() < 1e-10, - "30+40+30 should be 100.0, got {}", - value - ); - } - - #[test] - fn test_merge_separate_groups_no_contamination() { - let bytes_a = serialize_accumulator_arroyo(&SumAccumulator::with_sum(100.0)); - let bytes_b = serialize_accumulator_arroyo(&SumAccumulator::with_sum(200.0)); - let batch = make_batch(vec![("host-a", bytes_a), ("host-b", bytes_b)]); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 2, "Two different hosts should stay separate"); - - // Verify values - for (key, (_, bytes)) in &groups { - let restored = - crate::engines::physical::accumulator_serde::deserialize_single_subpopulation( - bytes, - &SketchType::Sum, - ) - .unwrap(); - let value = restored - .query(promql_utilities::query_logics::enums::Statistic::Sum, None) - .unwrap(); - if key[0] == "host-a" { - assert!((value - 100.0).abs() < 1e-10); - } else { - assert!((value - 200.0).abs() < 1e-10); - } - } - } - - #[test] - fn test_merge_kll_sketches() { - let mut kll1 = DatasketchesKLLAccumulator::new(200); - kll1.update(1.0); - kll1.update(2.0); - let mut kll2 = DatasketchesKLLAccumulator::new(200); - kll2.update(3.0); - kll2.update(4.0); - - let bytes1 = serialize_accumulator_arroyo(&kll1); - let bytes2 = serialize_accumulator_arroyo(&kll2); - let batch = make_batch(vec![("host-a", bytes1), ("host-a", bytes2)]); - - let logical = make_logical_node(SketchType::KLL); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - // Verify merged KLL has data from both - let (_, merged_bytes) = groups.values().next().unwrap(); - let restored = - crate::engines::physical::accumulator_serde::deserialize_single_subpopulation( - merged_bytes, - &SketchType::KLL, - ) - .unwrap(); - let mut kwargs = std::collections::HashMap::new(); - kwargs.insert("quantile".to_string(), "1.0".to_string()); - let max = restored - .query( - promql_utilities::query_logics::enums::Statistic::Quantile, - Some(&kwargs), - ) - .unwrap(); - assert!( - (max - 4.0).abs() < 1e-10, - "Max quantile should be 4.0, got {}", - max - ); - } - - #[test] - fn test_merge_set_aggregators() { - let mut set1 = SetAggregatorAccumulator::new(); - set1.add_key(KeyByLabelValues { - labels: vec!["a".to_string()], - }); - let mut set2 = SetAggregatorAccumulator::new(); - set2.add_key(KeyByLabelValues { - labels: vec!["b".to_string()], - }); - - let bytes1 = serialize_accumulator_arroyo(&set1); - let bytes2 = serialize_accumulator_arroyo(&set2); - let batch = make_batch(vec![("host-a", bytes1), ("host-a", bytes2)]); - - let logical = make_logical_node(SketchType::SetAggregator); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - let (_, merged_bytes) = groups.values().next().unwrap(); - let restored = deserialize_accumulator(merged_bytes, &SketchType::SetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 2, "Union of two sets should have 2 keys"); - } - - #[test] - fn test_merge_empty_batch() { - let schema = Arc::new(Schema::new(vec![ - Field::new("host", DataType::Utf8, true), - Field::new("sketch", DataType::Binary, false), - ])); - let host_array = StringArray::from(Vec::<&str>::new()); - let sketch_array = arrow::array::BinaryArray::from(Vec::<&[u8]>::new()); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(host_array) as ArrayRef, - Arc::new(sketch_array) as ArrayRef, - ], - ) - .unwrap(); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 0, "Empty batch should produce 0 groups"); - } -} diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 65415d6b..71fe6ced 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -820,167 +820,6 @@ impl SimpleEngine { Ok(results) } - /// Execute a query using the plan-based approach (for testing) - /// - /// This is an alternative execution path that uses DataFusion logical/physical - /// plans instead of the existing execute_query_pipeline. - /// - /// # Arguments - /// * `context` - The query execution context - /// - /// # Returns - /// A Result containing the query results or an error - #[allow(dead_code)] - pub async fn execute_plan( - &self, - context: &QueryExecutionContext, - ) -> Result, String> { - use datafusion::execution::context::SessionContext; - use datafusion::physical_plan::collect; - - use super::physical::conversion::record_batch_to_result_map; - - let total_start = Instant::now(); - - // 1. Build logical plan from context - let plan_build_start = Instant::now(); - let logical_plan = context - .to_logical_plan() - .map_err(|e| format!("Failed to build logical plan: {}", e))?; - debug!( - "[LATENCY] DataFusion: logical plan build: {:.2}ms", - plan_build_start.elapsed().as_secs_f64() * 1000.0 - ); - debug!( - "DataFusion logical plan:\n{}", - logical_plan.display_indent() - ); - - // 2. Create session context with our custom extension planner - let physical_plan_start = Instant::now(); - let session_ctx = SessionContext::new(); - #[allow(deprecated)] - let state = session_ctx.state().with_query_planner(std::sync::Arc::new( - super::physical::CustomQueryPlanner::new(self.store.clone()), - )); - - // 3. Create physical plan - let physical_plan = state - .create_physical_plan(&logical_plan) - .await - .map_err(|e| format!("Failed to create physical plan: {}", e))?; - debug!( - "[LATENCY] DataFusion: physical plan creation: {:.2}ms", - physical_plan_start.elapsed().as_secs_f64() * 1000.0 - ); - - // 4. Execute - let execute_start = Instant::now(); - let task_ctx = session_ctx.task_ctx(); - let batches = collect(physical_plan, task_ctx) - .await - .map_err(|e| format!("Failed to execute plan: {}", e))?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - debug!( - "[LATENCY] DataFusion: plan execution: {:.2}ms, {} batch(es), {} total rows", - execute_start.elapsed().as_secs_f64() * 1000.0, - batches.len(), - total_rows - ); - - // 5. Convert results - let convert_start = Instant::now(); - let label_names: Vec<&str> = context - .metadata - .query_output_labels - .labels - .iter() - .map(String::as_str) - .collect(); - - let mut all_results: HashMap, f64> = HashMap::new(); - for batch in &batches { - let batch_results = record_batch_to_result_map(batch, &label_names, "value") - .map_err(|e| format!("Failed to convert results: {}", e))?; - all_results.extend(batch_results); - } - debug!( - "[LATENCY] DataFusion: result conversion: {:.2}ms, {} output rows", - convert_start.elapsed().as_secs_f64() * 1000.0, - all_results.len() - ); - - // 6. Format results - let format_start = Instant::now(); - let results = self.format_final_results( - all_results, - &context.metadata.statistic_to_compute, - &context.metric, - false, - ); - debug!( - "[LATENCY] DataFusion: result formatting: {:.2}ms, {} results", - format_start.elapsed().as_secs_f64() * 1000.0, - results.len() - ); - - debug!( - "[LATENCY] DataFusion: total execute_plan: {:.2}ms", - total_start.elapsed().as_secs_f64() * 1000.0 - ); - - Ok(results) - } - - /// Executes a pre-built DataFusion logical plan and returns results. - /// - /// This was the entry point for the DataFusion-based binary arithmetic - /// dispatch path, cut over to a native implementation in #567. Unlike its - /// sibling `execute_plan` (still called by DataFusion-path tests), this - /// function has zero callers anywhere in the repo, including tests — it - /// is genuinely dead code, kept only in case the native cutover needs to - /// be reverted. - #[allow(dead_code)] - pub async fn execute_logical_plan( - &self, - logical_plan: datafusion::logical_expr::LogicalPlan, - label_names: Vec, - metric: &str, - statistic: &Statistic, - ) -> Result, String> { - use datafusion::execution::context::SessionContext; - use datafusion::physical_plan::collect; - - use super::physical::conversion::record_batch_to_result_map; - - // Create session context with our custom extension planner - let session_ctx = SessionContext::new(); - #[allow(deprecated)] - let state = session_ctx.state().with_query_planner(std::sync::Arc::new( - super::physical::CustomQueryPlanner::new(self.store.clone()), - )); - - let physical_plan = state - .create_physical_plan(&logical_plan) - .await - .map_err(|e| format!("Failed to create physical plan: {}", e))?; - - let task_ctx = session_ctx.task_ctx(); - let batches = collect(physical_plan, task_ctx) - .await - .map_err(|e| format!("Failed to execute plan: {}", e))?; - - let label_name_strs: Vec<&str> = label_names.iter().map(String::as_str).collect(); - let mut all_results: HashMap, f64> = HashMap::new(); - for batch in &batches { - let batch_results = record_batch_to_result_map(batch, &label_name_strs, "value") - .map_err(|e| format!("Failed to convert results: {}", e))?; - all_results.extend(batch_results); - } - - Ok(self.format_final_results(all_results, statistic, metric, false)) - } - /// Formats unformatted results into final InstantVectorElement format. /// /// For top-k queries the rows are always sorted by value descending (that's diff --git a/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs b/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs deleted file mode 100644 index c90e8e56..00000000 --- a/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs +++ /dev/null @@ -1,339 +0,0 @@ -//! Accumulator Serde Round-Trip Tests -//! -//! Tests that exercise accumulator_serde.rs directly (no engine needed). -//! Verifies serialize -> deserialize round-trip for all accumulator types. - -#[cfg(test)] -mod tests { - use crate::data_model::SerializableToSink; - use crate::data_model::{KeyByLabelValues, Measurement}; - use crate::engines::physical::accumulator_serde::{ - deserialize_accumulator, deserialize_keys_accumulator, deserialize_multiple_subpopulation, - deserialize_single_subpopulation, serialize_accumulator_arroyo, - }; - use crate::precompute_operators::{ - CountMinSketchAccumulator, DatasketchesKLLAccumulator, DeltaSetAggregatorAccumulator, - HydraKllSketchAccumulator, IncreaseAccumulator, MultipleIncreaseAccumulator, - SetAggregatorAccumulator, SumAccumulator, - }; - use datafusion_summary_library::SketchType; - use promql_utilities::query_logics::enums::{AggregationType, Statistic}; - use std::collections::HashMap; - - // ======================================================================== - // Full round-trip tests (serialize_arroyo -> deserialize) - // ======================================================================== - - #[test] - fn test_round_trip_sum() { - 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(), AggregationType::Sum); - - // Query the restored accumulator via single subpopulation - let restored_single = deserialize_single_subpopulation(&bytes, &SketchType::Sum).unwrap(); - let value = restored_single.query(Statistic::Sum, None).unwrap(); - assert!((value - 42.5).abs() < 1e-10, "Expected 42.5, got {}", value); - } - - #[test] - fn test_round_trip_kll() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in [1.0, 2.0, 3.0, 4.0, 5.0] { - kll.update(v); - } - - let bytes = serialize_accumulator_arroyo(&kll); - let restored = deserialize_accumulator(&bytes, &SketchType::KLL).unwrap(); - assert_eq!( - restored.get_accumulator_type(), - AggregationType::DatasketchesKLL - ); - - // Query quantile via single subpopulation - let restored_single = deserialize_single_subpopulation(&bytes, &SketchType::KLL).unwrap(); - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.5".to_string()); - let median = restored_single - .query(Statistic::Quantile, Some(&kwargs)) - .unwrap(); - // Median of [1,2,3,4,5] should be ~3.0 - assert!( - (1.0..=5.0).contains(&median), - "Median should be in [1,5], got {}", - median - ); - } - - #[test] - fn test_round_trip_set_aggregator() { - let mut set_acc = SetAggregatorAccumulator::new(); - set_acc.add_key(KeyByLabelValues { - labels: vec!["web".to_string()], - }); - set_acc.add_key(KeyByLabelValues { - labels: vec!["api".to_string()], - }); - set_acc.add_key(KeyByLabelValues { - labels: vec!["worker".to_string()], - }); - - let bytes = serialize_accumulator_arroyo(&set_acc); - let restored = deserialize_accumulator(&bytes, &SketchType::SetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 3, "Expected 3 keys, got {}", keys.len()); - } - - #[test] - fn test_round_trip_delta_set_aggregator() { - let mut delta_set = DeltaSetAggregatorAccumulator::new(); - delta_set.add_key(KeyByLabelValues { - labels: vec!["endpoint-a".to_string()], - }); - delta_set.add_key(KeyByLabelValues { - labels: vec!["endpoint-b".to_string()], - }); - - let bytes = serialize_accumulator_arroyo(&delta_set); - let restored = deserialize_accumulator(&bytes, &SketchType::DeltaSetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 2, "Expected 2 keys, got {}", keys.len()); - } - - #[test] - fn test_round_trip_multiple_increase() { - let key1 = KeyByLabelValues { - labels: vec!["web".to_string()], - }; - let key2 = KeyByLabelValues { - labels: vec!["api".to_string()], - }; - let mut increases = HashMap::new(); - increases.insert( - key1.clone(), - IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(100.0), 10), - ); - increases.insert( - key2.clone(), - IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(200.0), 10), - ); - - let acc = MultipleIncreaseAccumulator::new_with_increases(increases); - let bytes = serialize_accumulator_arroyo(&acc); - - let restored = deserialize_accumulator(&bytes, &SketchType::MultipleIncrease).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 2, "Expected 2 keys, got {}", keys.len()); - - // Query via multiple subpopulation - let restored_multi = - deserialize_multiple_subpopulation(&bytes, &SketchType::MultipleIncrease).unwrap(); - let val = restored_multi - .query(Statistic::Increase, &key1, None) - .unwrap(); - assert!( - (val - 100.0).abs() < 1e-10, - "Expected increase=100.0 for key1, got {}", - val - ); - } - - #[test] - fn test_round_trip_hydra_kll() { - // HydraKLL: serialize_to_bytes IS MessagePack, so serialize_accumulator_arroyo - // falls through to serialize_to_bytes - let mut hydra = HydraKllSketchAccumulator::new(1, 1, 200); - // Use the public update method with a key - hydra.update( - &KeyByLabelValues { - labels: vec!["sub-key".to_string()], - }, - 42.0, - ); - - let bytes = serialize_accumulator_arroyo(&hydra); - let restored = deserialize_accumulator(&bytes, &SketchType::HydraKLL).unwrap(); - assert_eq!(restored.get_accumulator_type(), AggregationType::HydraKLL); - } - - #[test] - fn test_round_trip_count_min_sketch() { - // CountMinSketch: serialize_to_bytes IS MessagePack - // Supported by both deserialize_accumulator (for merging) and - // deserialize_multiple_subpopulation (for querying by sub-key) - let cms = CountMinSketchAccumulator::new(2, 3); - - let bytes = serialize_accumulator_arroyo(&cms); - - let restored = deserialize_accumulator(&bytes, &SketchType::CountMinSketch).unwrap(); - 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(), - AggregationType::CountMinSketch - ); - } - - // ======================================================================== - // Keys accumulator round-trip - // ======================================================================== - - #[test] - fn test_deserialize_keys_delta_set() { - let mut delta_set = DeltaSetAggregatorAccumulator::new(); - delta_set.add_key(KeyByLabelValues { - labels: vec!["key-a".to_string()], - }); - let bytes = serialize_accumulator_arroyo(&delta_set); - - let restored = - deserialize_keys_accumulator(&bytes, &SketchType::DeltaSetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 1); - } - - #[test] - fn test_deserialize_keys_set_aggregator() { - let mut set_acc = SetAggregatorAccumulator::new(); - set_acc.add_key(KeyByLabelValues { - labels: vec!["k1".to_string()], - }); - let bytes = serialize_accumulator_arroyo(&set_acc); - - let restored = deserialize_keys_accumulator(&bytes, &SketchType::SetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 1); - } - - // ======================================================================== - // Not-implemented types - // ======================================================================== - - #[test] - fn test_not_implemented_increase() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_accumulator(&bytes, &SketchType::Increase); - assert!(result.is_err()); - let err_msg = match result { - Err(e) => format!("{}", e), - Ok(_) => panic!("Expected error"), - }; - assert!( - err_msg.contains("Not") || err_msg.contains("not"), - "Expected NotImplemented error, got: {}", - err_msg - ); - } - - #[test] - fn test_not_implemented_minmax() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_accumulator(&bytes, &SketchType::MinMax); - assert!(result.is_err()); - } - - // ======================================================================== - // Error handling tests - // ======================================================================== - - #[test] - fn test_corrupted_bytes_sum() { - // MessagePack can decode some short byte sequences as valid integers/floats - // (e.g., 0xFF = -1 as negative fixint). Use 0xCB (float64 marker) followed - // by insufficient bytes to force a decode error. - let garbage = vec![0xCB, 0x01, 0x02]; - let result = deserialize_accumulator(&garbage, &SketchType::Sum); - assert!( - result.is_err(), - "Corrupted bytes should produce an error, not a panic" - ); - } - - #[test] - fn test_corrupted_bytes_kll() { - let garbage = vec![0xFF, 0xFE, 0xFD, 0xFC]; - let result = deserialize_accumulator(&garbage, &SketchType::KLL); - assert!( - result.is_err(), - "Corrupted bytes should produce an error, not a panic" - ); - } - - #[test] - fn test_empty_bytes_sum() { - let result = deserialize_accumulator(&[], &SketchType::Sum); - assert!(result.is_err(), "Empty bytes should produce an error"); - } - - #[test] - fn test_empty_bytes_kll() { - let result = deserialize_accumulator(&[], &SketchType::KLL); - assert!(result.is_err(), "Empty bytes should produce an error"); - } - - #[test] - fn test_empty_bytes_set_aggregator() { - let result = deserialize_accumulator(&[], &SketchType::SetAggregator); - assert!(result.is_err(), "Empty bytes should produce an error"); - } - - #[test] - fn test_empty_bytes_delta_set() { - let result = deserialize_accumulator(&[], &SketchType::DeltaSetAggregator); - assert!(result.is_err(), "Empty bytes should produce an error"); - } - - // ======================================================================== - // Serialize dispatch verification - // ======================================================================== - - #[test] - fn test_serialize_arroyo_dispatch_sum_uses_arroyo_path() { - // SumAccumulator has a separate arroyo format (MessagePack f64) - // while its native serialize_to_bytes uses little-endian f64 - let acc = SumAccumulator::with_sum(42.0); - let arroyo_bytes = serialize_accumulator_arroyo(&acc); - let native_bytes = acc.serialize_to_bytes(); - - // They should differ because arroyo uses MessagePack - assert_ne!( - arroyo_bytes, native_bytes, - "SumAccumulator arroyo and native serialization should differ" - ); - - // Verify the arroyo bytes can be deserialized - let restored = deserialize_accumulator(&arroyo_bytes, &SketchType::Sum).unwrap(); - assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); - } - - #[test] - fn test_serialize_arroyo_dispatch_kll_uses_native() { - // KLL's serialize_to_bytes already uses MessagePack, so arroyo falls through - let mut kll = DatasketchesKLLAccumulator::new(200); - kll.update(1.0); - let arroyo_bytes = serialize_accumulator_arroyo(&kll); - let native_bytes = kll.serialize_to_bytes(); - - // They should be the same since KLL's native IS MessagePack - assert_eq!( - arroyo_bytes, native_bytes, - "KLL arroyo and native serialization should be the same" - ); - } - - #[test] - fn test_unsupported_keys_type() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_keys_accumulator(&bytes, &SketchType::Sum); - assert!( - result.is_err(), - "Sum should not be supported as a keys accumulator" - ); - } -} diff --git a/asap-query-engine/src/tests/datafusion/mod.rs b/asap-query-engine/src/tests/datafusion/mod.rs deleted file mode 100644 index fd946643..00000000 --- a/asap-query-engine/src/tests/datafusion/mod.rs +++ /dev/null @@ -1,14 +0,0 @@ -//! DataFusion execution path tests. -//! -//! Tests for the logical plan builder, physical execution operators, -//! and accumulator serialization that back the DataFusion-based query path. - -pub mod accumulator_serde_tests; -pub mod dispatch_arithmetic_tests; -pub mod plan_builder_binary_tests; -pub mod plan_builder_regression_tests; -pub mod plan_execution_dual_input_tests; -pub mod plan_execution_temporal_tests; -pub mod plan_execution_tests; -pub mod range_query_arithmetic_tests; -pub mod structural_matching_tests; 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 deleted file mode 100644 index c6639069..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs +++ /dev/null @@ -1,190 +0,0 @@ -//! Binary plan builder tests. -//! -//! Tests that `build_binary_vector_plan` and `build_scalar_plan` produce correct -//! DataFusion logical plan structures. - -#[cfg(test)] -mod tests { - use crate::data_model::AggregationIdInfo; - use crate::engines::logical::plan_builder::{build_binary_vector_plan, build_scalar_plan}; - use crate::engines::simple_engine::{ - QueryExecutionContext, QueryMetadata, StoreQueryParams, StoreQueryPlan, - }; - 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::{AggregationType, Statistic}; - use std::collections::HashMap; - - fn make_context( - metric: &str, - statistic: Statistic, - labels: Vec<&str>, - ) -> QueryExecutionContext { - let label_strings: Vec = labels.into_iter().map(String::from).collect(); - QueryExecutionContext { - metric: metric.to_string(), - metadata: QueryMetadata { - query_output_labels: KeyByLabelNames::new(label_strings.clone()), - statistic_to_compute: statistic, - query_kwargs: HashMap::new(), - }, - store_plan: StoreQueryPlan { - values_query: StoreQueryParams { - metric: metric.to_string(), - aggregation_id: 1, - start_timestamp: 1000, - end_timestamp: 2000, - is_exact_query: true, - }, - keys_query: None, - }, - agg_info: AggregationIdInfo { - aggregation_id_for_key: 1, - aggregation_id_for_value: 1, - aggregation_type_for_key: AggregationType::Sum, - aggregation_type_for_value: AggregationType::Sum, - }, - do_merge: false, - spatial_filter: String::new(), - query_time: 2000, - grouping_labels: KeyByLabelNames::new(label_strings.clone()), - aggregated_labels: KeyByLabelNames::empty(), - } - } - - fn collect_node_names(plan: &LogicalPlan) -> Vec { - let mut names = Vec::new(); - collect_recursive(plan, &mut names); - names - } - - fn collect_recursive(plan: &LogicalPlan, names: &mut Vec) { - match plan { - LogicalPlan::Extension(ext) => { - names.push(ext.node.name().to_string()); - for input in ext.node.inputs() { - collect_recursive(input, names); - } - } - LogicalPlan::Projection(p) => { - names.push("Projection".to_string()); - collect_recursive(&p.input, names); - } - LogicalPlan::Join(j) => { - names.push("Join".to_string()); - collect_recursive(&j.left, names); - collect_recursive(&j.right, names); - } - LogicalPlan::SubqueryAlias(a) => { - names.push("SubqueryAlias".to_string()); - collect_recursive(&a.input, names); - } - other => { - names.push( - format!("{:?}", other) - .split('(') - .next() - .unwrap_or("Unknown") - .to_string(), - ); - } - } - } - - fn contains_node(plan: &LogicalPlan, name: &str) -> bool { - collect_node_names(plan).iter().any(|n| n == name) - } - - #[test] - fn test_binary_vector_plan_structure_divide() { - let lhs_ctx = make_context("errors", Statistic::Sum, vec!["host"]); - let rhs_ctx = make_context("requests", Statistic::Sum, vec!["host"]); - let lhs_plan = lhs_ctx.to_logical_plan().unwrap(); - let rhs_plan = rhs_ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(T_DIV); - let plan = - build_binary_vector_plan(lhs_plan, rhs_plan, &op, vec!["host".to_string()]).unwrap(); - - let names = collect_node_names(&plan); - assert_eq!(names[0], "Projection", "Root should be Projection"); - assert!(contains_node(&plan, "Join"), "Plan should contain a Join"); - let alias_count = names.iter().filter(|n| *n == "SubqueryAlias").count(); - assert_eq!( - alias_count, 2, - "Plan should contain two SubqueryAlias nodes" - ); - } - - #[test] - fn test_binary_vector_plan_all_operators() { - let ops = [T_ADD, T_SUB, T_MUL, T_DIV, T_POW, T_MOD]; - for op_id in ops { - let lhs_ctx = make_context("metric_a", Statistic::Sum, vec!["host"]); - let rhs_ctx = make_context("metric_b", Statistic::Sum, vec!["host"]); - let lhs_plan = lhs_ctx.to_logical_plan().unwrap(); - let rhs_plan = rhs_ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(op_id); - let result = - build_binary_vector_plan(lhs_plan, rhs_plan, &op, vec!["host".to_string()]); - assert!( - result.is_ok(), - "Operator {:?} should produce a valid plan", - op - ); - let names = collect_node_names(&result.unwrap()); - assert_eq!(names[0], "Projection"); - } - } - - #[test] - fn test_scalar_right_plan_structure() { - let ctx = make_context("errors", Statistic::Sum, vec!["host"]); - let vector_plan = ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(T_MUL); - let plan = - build_scalar_plan(vector_plan, 100.0, &op, false, vec!["host".to_string()]).unwrap(); - - let names = collect_node_names(&plan); - assert_eq!(names[0], "Projection", "Root should be Projection"); - assert!( - !contains_node(&plan, "Join"), - "Scalar plan should not have a Join" - ); - assert!( - !contains_node(&plan, "SubqueryAlias"), - "Scalar plan should not have SubqueryAlias" - ); - } - - #[test] - fn test_scalar_left_plan_structure() { - let ctx = make_context("success", Statistic::Sum, vec!["host"]); - let vector_plan = ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(T_SUB); - let plan = - build_scalar_plan(vector_plan, 1.0, &op, true, vec!["host".to_string()]).unwrap(); - - let names = collect_node_names(&plan); - assert_eq!(names[0], "Projection"); - assert!(!contains_node(&plan, "Join")); - } - - #[test] - fn test_scalar_left_division_plan_structure() { - // 1.0 / rate(metric[5m]) — scalar on left with Div - let ctx = make_context("metric", Statistic::Sum, vec!["host"]); - let vector_plan = ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(T_DIV); - let result = build_scalar_plan(vector_plan, 1.0, &op, true, vec!["host".to_string()]); - assert!( - result.is_ok(), - "scalar-left division plan should build without error" - ); - } -} 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 deleted file mode 100644 index 8764a806..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs +++ /dev/null @@ -1,210 +0,0 @@ -//! Plan Builder Regression Tests -//! -//! Tests covering gaps in the existing plan_builder.rs inline tests: -//! all Statistic variants, kwargs propagation, error paths, edge cases. - -#[cfg(test)] -mod tests { - use crate::data_model::AggregationIdInfo; - use crate::engines::simple_engine::{ - QueryExecutionContext, QueryMetadata, StoreQueryParams, StoreQueryPlan, - }; - use promql_utilities::data_model::KeyByLabelNames; - use promql_utilities::query_logics::enums::{AggregationType, Statistic}; - use std::collections::HashMap; - - fn create_context( - statistic: Statistic, - aggregation_type: AggregationType, - output_labels: Vec<&str>, - kwargs: HashMap, - ) -> QueryExecutionContext { - let output_labels_vec: Vec = output_labels.into_iter().map(String::from).collect(); - let labels = KeyByLabelNames { - labels: output_labels_vec, - }; - QueryExecutionContext { - metric: "test_metric".to_string(), - metadata: QueryMetadata { - query_output_labels: labels.clone(), - statistic_to_compute: statistic, - query_kwargs: kwargs, - }, - store_plan: StoreQueryPlan { - values_query: StoreQueryParams { - metric: "test_metric".to_string(), - aggregation_id: 1, - start_timestamp: 1000, - end_timestamp: 2000, - is_exact_query: true, - }, - keys_query: None, - }, - agg_info: AggregationIdInfo { - aggregation_id_for_key: 1, - aggregation_id_for_value: 1, - aggregation_type_for_key: AggregationType::Sum, - aggregation_type_for_value: aggregation_type, - }, - do_merge: false, - spatial_filter: String::new(), - query_time: 2000, - grouping_labels: labels, - aggregated_labels: KeyByLabelNames::empty(), - } - } - - // ======================================================================== - // All Statistic variants map without panic - // ======================================================================== - - #[test] - fn test_all_statistics_map_without_panic() { - let statistics = vec![ - (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 { - let ctx = create_context(stat, agg_type, vec!["host"], HashMap::new()); - let result = ctx.map_statistic_to_infer_operation(); - assert!( - result.is_ok(), - "Statistic {:?} should map successfully, got: {:?}", - stat, - result.err() - ); - } - } - - // ======================================================================== - // TopK kwargs propagation - // ======================================================================== - - #[test] - fn test_topk_kwargs_propagate() { - use datafusion_summary_library::InferOperation; - let mut kwargs = HashMap::new(); - kwargs.insert("k".to_string(), "5".to_string()); - - let ctx = create_context( - Statistic::Topk, - AggregationType::CountMinSketch, - vec!["host"], - kwargs, - ); - match ctx.map_statistic_to_infer_operation().unwrap() { - InferOperation::TopK(k) => assert_eq!(k, 5, "Expected k=5, got {}", k), - other => panic!("Expected TopK, got {:?}", other), - } - } - - #[test] - fn test_topk_default_k() { - use datafusion_summary_library::InferOperation; - let ctx = create_context( - Statistic::Topk, - AggregationType::CountMinSketch, - vec!["host"], - HashMap::new(), - ); - match ctx.map_statistic_to_infer_operation().unwrap() { - InferOperation::TopK(k) => assert_eq!(k, 10, "Default k should be 10, got {}", k), - other => panic!("Expected TopK, got {:?}", other), - } - } - - // ======================================================================== - // Statistic-to-operation mapping - // ======================================================================== - - #[test] - fn test_cardinality_to_count_distinct() { - use datafusion_summary_library::InferOperation; - let ctx = create_context( - Statistic::Cardinality, - AggregationType::SetAggregator, - vec!["host"], - HashMap::new(), - ); - assert!(matches!( - ctx.map_statistic_to_infer_operation().unwrap(), - InferOperation::CountDistinct - )); - } - - #[test] - fn test_rate_to_extract_rate() { - use datafusion_summary_library::InferOperation; - let ctx = create_context( - Statistic::Rate, - AggregationType::Increase, - vec!["host"], - HashMap::new(), - ); - assert!(matches!( - ctx.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractRate - )); - } - - #[test] - fn test_count_to_extract_count() { - use datafusion_summary_library::InferOperation; - let ctx = create_context( - Statistic::Count, - AggregationType::Sum, - vec!["host"], - HashMap::new(), - ); - assert!(matches!( - ctx.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractCount - )); - } - - // ======================================================================== - // Error paths - // ======================================================================== - - #[test] - fn test_unknown_agg_type_errors() { - // SingleSubpopulation is a legacy wrapper variant not mapped to a SketchType - let ctx = create_context( - Statistic::Sum, - AggregationType::SingleSubpopulation, - vec!["host"], - HashMap::new(), - ); - let result = ctx.to_logical_plan(); - assert!(result.is_err(), "Unmapped aggregation type should error"); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains("Unknown"), - "Error should mention Unknown, got: {}", - err_msg - ); - } - - // ======================================================================== - // Edge cases - // ======================================================================== - - #[test] - fn test_plan_with_empty_labels() { - let ctx = create_context(Statistic::Sum, AggregationType::Sum, vec![], HashMap::new()); - let result = ctx.to_logical_plan(); - assert!( - result.is_ok(), - "Empty labels should still build a plan: {:?}", - result.err() - ); - } -} 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 deleted file mode 100644 index 649b2b24..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs +++ /dev/null @@ -1,312 +0,0 @@ -//! Plan Execution Multi-Population Tests -//! -//! Tests for multi-population accumulator execution: -//! -//! **Self-keyed (single-input):** MultipleIncrease, MultipleSum, MultipleMinMax -//! These carry their own keys via get_keys(). No separate keys stream needed. -//! grouping_labels = [], aggregated_labels = [all output labels] -//! -//! **Dual-input:** HydraKLL, CountMinSketch -//! These need a separate DeltaSetAggregator keys stream to enumerate sub-keys. - -#[cfg(test)] -mod tests { - use crate::data_model::{AggregationType, KeyByLabelValues, Measurement}; - use crate::precompute_operators::{ - CountMinSketchAccumulator, DeltaSetAggregatorAccumulator, HydraKllSketchAccumulator, - IncreaseAccumulator, MultipleIncreaseAccumulator, - }; - use crate::tests::test_utilities::engine_factories::*; - use std::collections::HashMap; - - /// Helper: create a MultipleIncreaseAccumulator with given sub-keys. - /// Each key is a Vec of label values (matching aggregated_labels order). - fn make_multi_increase( - keys_and_values: Vec<(Vec<&str>, f64, f64)>, - ) -> MultipleIncreaseAccumulator { - let mut increases = HashMap::new(); - for (key_labels, start_val, end_val) in keys_and_values { - let key = KeyByLabelValues { - labels: key_labels.iter().map(|s| s.to_string()).collect(), - }; - increases.insert( - key, - IncreaseAccumulator::new( - Measurement::new(start_val), - 0, - Measurement::new(end_val), - 10, - ), - ); - } - MultipleIncreaseAccumulator::new_with_increases(increases) - } - - // ======================================================================== - // Self-keyed: MultipleIncrease (single-input, keys from accumulator) - // MultipleIncrease is a collection of Increase accumulators. - // grouping_labels = [], aggregated_labels = [host, endpoint] - // Query is simply increase(metric[window]). - // ======================================================================== - - #[tokio::test] - async fn test_self_keyed_multiple_increase() { - // Sub-keys: (host-a, endpoint-1) increase=100, (host-a, endpoint-2) increase=200 - let acc = make_multi_increase(vec![ - (vec!["host-a", "endpoint-1"], 0.0, 100.0), - (vec!["host-a", "endpoint-2"], 0.0, 200.0), - ]); - - let engine = create_engine_single_pop_with_aggregated( - "http_requests_total", - AggregationType::MultipleIncrease, - vec![], - vec!["host", "endpoint"], - vec![(None, Box::new(acc))], - "increase(http_requests_total[10s])", - ); - - let results = execute_new_plan(&engine, "increase(http_requests_total[10s])", 1000.0).await; - - assert!( - !results.is_empty(), - "Self-keyed multi-pop should produce results; got 0" - ); - - // Should have 2 results (one per sub-key) - assert_eq!( - results.len(), - 2, - "Expected 2 results (2 sub-keys), got {}", - results.len() - ); - } - - #[tokio::test] - async fn test_self_keyed_multiple_increase_single_key() { - let acc = make_multi_increase(vec![(vec!["host-a", "svc-web"], 0.0, 50.0)]); - - let engine = create_engine_single_pop_with_aggregated( - "http_requests_total", - AggregationType::MultipleIncrease, - vec![], - vec!["host", "service"], - vec![(None, Box::new(acc))], - "increase(http_requests_total[10s])", - ); - - let results = execute_new_plan(&engine, "increase(http_requests_total[10s])", 1000.0).await; - - assert_eq!( - results.len(), - 1, - "Expected 1 result (1 sub-key), got {}", - results.len() - ); - } - - // ======================================================================== - // HydraKLL + DeltaSetAggregator (quantile dual-input) - // ======================================================================== - - #[tokio::test] - async fn test_dual_hydra_kll_delta_set() { - // HydraKLL: single accumulator, no spatial GROUP BY. - // Sub-keys are ["host", "endpoint"] — tracked by DeltaSet, queryable in HydraKLL. - // 2 columns per sub-key (host, endpoint). - let mut hydra = HydraKllSketchAccumulator::new(1, 2, 200); - hydra.update( - &KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-a".to_string()], - }, - 10.0, - ); - hydra.update( - &KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-a".to_string()], - }, - 20.0, - ); - hydra.update( - &KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-b".to_string()], - }, - 100.0, - ); - - // DeltaSet enumerates which sub-keys exist - let mut keys = DeltaSetAggregatorAccumulator::new(); - keys.add_key(KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-a".to_string()], - }); - keys.add_key(KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-b".to_string()], - }); - - let engine = create_engine_dual_input( - "request_duration", - 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 - vec![(None, Box::new(keys))], // store key = None - "quantile(0.5, request_duration) by (host, endpoint)", - ); - - let results = execute_new_plan( - &engine, - "quantile(0.5, request_duration) by (host, endpoint)", - 1000.0, - ) - .await; - - // Should have 2 results: one per (host, endpoint) sub-key - assert_eq!( - results.len(), - 2, - "HydraKLL dual-input should produce 2 results, got {}", - results.len() - ); - } - - // ======================================================================== - // CountMinSketch + DeltaSetAggregator (frequency dual-input) - // ======================================================================== - - #[tokio::test] - async fn test_dual_count_min_delta_set() { - // CountMinSketch: single accumulator, no spatial GROUP BY. - // Sub-keys are ["host", "event"] — tracked by DeltaSet, queryable in CMS. - let cms = CountMinSketchAccumulator::new(2, 3); - - let mut keys = DeltaSetAggregatorAccumulator::new(); - keys.add_key(KeyByLabelValues { - labels: vec!["host-a".to_string(), "evt-1".to_string()], - }); - - let engine = create_engine_dual_input( - "event_frequency", - 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 - vec![(None, Box::new(keys))], // store key = None - "count(event_frequency) by (host, event)", - ); - - let results = - execute_new_plan(&engine, "count(event_frequency) by (host, event)", 1000.0).await; - - // CMS query may return 0 for un-updated keys, but should not error - assert!(!results.is_empty(), "CMS dual-input should produce results"); - } - - // ======================================================================== - // Self-keyed: multiple accumulators (multiple store entries) - // ======================================================================== - - #[tokio::test] - async fn test_self_keyed_multiple_accumulators() { - // Two separate MultipleIncrease accumulators in the store (both with key=None) - // acc1 has (host-a, ep-1) and (host-a, ep-2) - // acc2 has (host-b, ep-3) - let acc1 = make_multi_increase(vec![ - (vec!["host-a", "ep-1"], 0.0, 10.0), - (vec!["host-a", "ep-2"], 0.0, 20.0), - ]); - let acc2 = make_multi_increase(vec![(vec!["host-b", "ep-3"], 0.0, 30.0)]); - - let engine = create_engine_single_pop_with_aggregated( - "requests", - AggregationType::MultipleIncrease, - vec![], - vec!["host", "endpoint"], - vec![(None, Box::new(acc1)), (None, Box::new(acc2))], - "increase(requests[10s])", - ); - - let results = execute_new_plan(&engine, "increase(requests[10s])", 1000.0).await; - - // 2 from acc1 + 1 from acc2 = 3 total - assert_eq!( - results.len(), - 3, - "Expected 3 results (2 + 1), got {}", - results.len() - ); - } - - // ======================================================================== - // Self-keyed: empty accumulator - // ======================================================================== - - #[tokio::test] - async fn test_self_keyed_empty_keys() { - // MultipleIncrease with no sub-keys -> 0 results - let empty = MultipleIncreaseAccumulator::new_with_increases(HashMap::new()); - - let engine = create_engine_single_pop_with_aggregated( - "requests", - AggregationType::MultipleIncrease, - vec![], - vec!["host", "endpoint"], - vec![(None, Box::new(empty))], - "increase(requests[10s])", - ); - - let results = execute_new_plan(&engine, "increase(requests[10s])", 1000.0).await; - - assert!( - results.is_empty(), - "Empty MultipleIncrease should give 0 results, got {}", - results.len() - ); - } - - // ======================================================================== - // Dual-input: mismatched spatial groups (HydraKLL) - // ======================================================================== - - #[tokio::test] - async fn test_dual_mismatched_spatial_groups() { - // No spatial GROUP BY: single HydraKLL and single DeltaSet. - // HydraKLL has data for (host-a, ep-1), DeltaSet tracks (host-b, ep-2). - // The DeltaSet key doesn't match any data in HydraKLL, so query returns 0/default. - let mut hydra = HydraKllSketchAccumulator::new(1, 2, 200); - hydra.update( - &KeyByLabelValues { - labels: vec!["host-a".to_string(), "ep-1".to_string()], - }, - 42.0, - ); - - let mut keys = DeltaSetAggregatorAccumulator::new(); - keys.add_key(KeyByLabelValues { - labels: vec!["host-b".to_string(), "ep-2".to_string()], - }); - - let engine = create_engine_dual_input( - "request_duration", - AggregationType::HydraKLL, - AggregationType::DeltaSetAggregator, - vec![], // no store GROUP BY - vec!["host", "endpoint"], // aggregated_labels - vec![(None, Box::new(hydra))], - vec![(None, Box::new(keys))], - "quantile(0.5, request_duration) by (host, endpoint)", - ); - - let results = execute_new_plan( - &engine, - "quantile(0.5, request_duration) by (host, endpoint)", - 1000.0, - ) - .await; - - // DeltaSet enumerates (host-b, ep-2) but HydraKLL has no data for that key. - // We just verify it doesn't panic. - let _ = results; - } -} 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 deleted file mode 100644 index fd081bd3..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs +++ /dev/null @@ -1,604 +0,0 @@ -//! Plan Execution Integration Tests — Temporal & Collapsable Queries -//! -//! Tests that verify the DataFusion plan-based execution path produces correct -//! results for temporal queries (sum_over_time, quantile_over_time) -//! and collapsable queries (spatial of temporal, e.g. sum by () (sum_over_time(...))). - -use crate::precompute_operators::sum_accumulator::SumAccumulator; -use std::collections::HashMap; - -#[cfg(test)] -mod tests { - use super::*; - use crate::data_model::{AggregationType, WindowType}; - use crate::precompute_operators::DatasketchesKLLAccumulator; - use crate::tests::test_utilities::engine_factories::*; - - type TemporalData = Vec<(u64, Option>, Box)>; - - // ======================================================================== - // Helper: build temporal data at 5 timestamps within a [5s] window - // Timestamps: 996_000, 997_000, 998_000, 999_000, 1_000_000 - // Query time: 1000.0 (= 1_000_000 ms) - // ======================================================================== - - const QUERY_TIME: f64 = 1000.0; - const TEMPORAL_TIMESTAMPS: [u64; 5] = [996_000, 997_000, 998_000, 999_000, 1_000_000]; - - // ======================================================================== - // OnlyTemporal Tests - // ======================================================================== - - #[tokio::test] - async fn test_temporal_sum_over_time_merges_across_timestamps() { - // Insert SumAccumulator data at 5 timestamps for one label group. - // sum_over_time should merge (sum) all values across timestamps. - let data: TemporalData = TEMPORAL_TIMESTAMPS - .iter() - .map(|&ts| { - ( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - ) - }) - .collect(); - - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5_000, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 1, "Expected 1 result for single group"); - // Merged sum: 10.0 * 5 timestamps = 50.0 - assert!( - (results[0].value - 50.0).abs() < 1e-10, - "Expected merged sum 50.0, got {}", - results[0].value - ); - } - - #[tokio::test] - async fn test_temporal_sum_over_time_single_timestamp() { - // Only 1 timestamp in range — should still work. - let data: TemporalData = vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(42.0)), - )]; - - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5_000, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 1); - assert!( - (results[0].value - 42.0).abs() < 1e-10, - "Expected 42.0, got {}", - results[0].value - ); - } - - #[tokio::test] - async fn test_temporal_sum_over_time_varying_values() { - // sum_over_time with different values at each timestamp. - // Verifies merge sums all values, not just takes latest. - let data: TemporalData = TEMPORAL_TIMESTAMPS - .iter() - .enumerate() - .map(|(i, &ts)| { - ( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum((i as f64 + 1.0) * 10.0)) - as Box, - ) - }) - .collect(); - - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5_000, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 1); - // 10 + 20 + 30 + 40 + 50 = 150 - assert!( - (results[0].value - 150.0).abs() < 1e-10, - "Expected merged sum 150.0, got {}", - results[0].value - ); - } - - #[tokio::test] - async fn test_temporal_quantile_over_time() { - // DatasketchesKLL at multiple timestamps, quantile_over_time(0.5, ...). - // Each KLL sketch has different values; merged sketch should give median. - let mut data: TemporalData = Vec::new(); - for (i, &ts) in TEMPORAL_TIMESTAMPS.iter().enumerate() { - let mut kll = DatasketchesKLLAccumulator::new(200); - // Insert values 10, 20, 30, 40, 50 at successive timestamps - kll.update((i as f64 + 1.0) * 10.0); - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(kll) as Box, - )); - } - - let query = "quantile_over_time(0.5, latency[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - data, - query, - 5_000, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 1); - // Median of {10, 20, 30, 40, 50} = 30.0 - assert!( - (results[0].value - 30.0).abs() < 5.0, - "Expected median ~30.0, got {}", - results[0].value - ); - } - - #[tokio::test] - async fn test_temporal_sum_over_time_multi_group() { - // Multiple label groups at multiple timestamps — verify per-group merging. - let mut data: TemporalData = Vec::new(); - for &ts in &TEMPORAL_TIMESTAMPS { - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )); - data.push(( - ts, - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(20.0)) as Box, - )); - } - - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5_000, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 2, "Expected 2 groups"); - - let result_map: HashMap = results - .iter() - .map(|r| (r.labels.labels[0].clone(), r.value)) - .collect(); - - assert!( - (result_map["host-a"] - 50.0).abs() < 1e-10, - "host-a: expected 50.0, got {}", - result_map["host-a"] - ); - assert!( - (result_map["host-b"] - 100.0).abs() < 1e-10, - "host-b: expected 100.0, got {}", - result_map["host-b"] - ); - } - - #[tokio::test] - async fn test_temporal_sum_over_time_empty_store() { - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![], - query, - 5_000, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert!( - results.is_empty(), - "Empty store should return 0 results, got {}", - results.len() - ); - } - - #[tokio::test] - async fn test_temporal_context_has_do_merge_true() { - // window_size_ms (1_000) must be smaller than the query's [5s] range so the - // derived do_merge (range_ms > window_size_ms) reflects a realistic config - // where the aggregation's bucket is finer than the requested range. - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(1.0)), - )], - query, - 1_000, - WindowType::Tumbling, - ); - - let context = engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .expect("Failed to build context"); - - assert!(context.do_merge, "Temporal queries must have do_merge=true"); - } - - #[tokio::test] - async fn test_temporal_context_has_correct_time_range() { - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(1.0)), - )], - query, - 5_000, - WindowType::Tumbling, - ); - - let context = engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .expect("Failed to build context"); - - // query_time = 1000.0 -> 1_000_000 ms - assert_eq!(context.query_time, 1_000_000); - // For [5s] range: start = end - 5000 = 995_000 - let start = context.store_plan.values_query.start_timestamp; - let end = context.store_plan.values_query.end_timestamp; - assert_eq!(end, 1_000_000, "End timestamp should be 1_000_000"); - assert_eq!( - start, 995_000, - "Start timestamp should be 995_000 for [5s] range" - ); - } - - // ======================================================================== - // Collapsable Tests (spatial of temporal) - // ======================================================================== - - #[tokio::test] - async fn test_collapsable_sum_of_sum_over_time() { - // sum by (host) (sum_over_time(metric[5s])) - // Multiple hosts at multiple timestamps. - let mut data: TemporalData = Vec::new(); - for &ts in &TEMPORAL_TIMESTAMPS { - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )); - data.push(( - ts, - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(20.0)) as Box, - )); - } - - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5_000, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 2, "Expected 2 groups (host-a, host-b)"); - - let result_map: HashMap = results - .iter() - .map(|r| (r.labels.labels[0].clone(), r.value)) - .collect(); - - // host-a: 10.0 * 5 timestamps = 50.0 - assert!( - (result_map["host-a"] - 50.0).abs() < 1e-10, - "host-a: expected 50.0, got {}", - result_map["host-a"] - ); - // host-b: 20.0 * 5 timestamps = 100.0 - assert!( - (result_map["host-b"] - 100.0).abs() < 1e-10, - "host-b: expected 100.0, got {}", - result_map["host-b"] - ); - } - - #[tokio::test] - async fn test_collapsable_sum_of_sum_over_time_varying_values() { - // sum by (host) (sum_over_time(metric[5s])) with varying values per timestamp - let mut data: TemporalData = Vec::new(); - for (i, &ts) in TEMPORAL_TIMESTAMPS.iter().enumerate() { - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum((i as f64 + 1.0) * 10.0)) - as Box, - )); - data.push(( - ts, - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum((i as f64 + 1.0) * 5.0)) - as Box, - )); - } - - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5_000, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 2, "Expected 2 groups"); - - let result_map: HashMap = results - .iter() - .map(|r| (r.labels.labels[0].clone(), r.value)) - .collect(); - - // host-a: 10 + 20 + 30 + 40 + 50 = 150.0 - assert!( - (result_map["host-a"] - 150.0).abs() < 1e-10, - "host-a: expected 150.0, got {}", - result_map["host-a"] - ); - // host-b: 5 + 10 + 15 + 20 + 25 = 75.0 - assert!( - (result_map["host-b"] - 75.0).abs() < 1e-10, - "host-b: expected 75.0, got {}", - result_map["host-b"] - ); - } - - #[tokio::test] - async fn test_collapsable_context_has_do_merge_true() { - // window_size_ms (1_000) must be smaller than the query's [5s] range so the - // derived do_merge (range_ms > window_size_ms) reflects a realistic config - // where the aggregation's bucket is finer than the requested range. - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(1.0)), - )], - query, - 1_000, - WindowType::Tumbling, - ); - - let context = engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .expect("Failed to build context"); - - assert!( - context.do_merge, - "Collapsable (OneTemporalOneSpatial) queries must have do_merge=true" - ); - } - - #[tokio::test] - async fn test_collapsable_output_labels_are_spatial() { - // Verify output labels are the spatial GROUP BY labels (host), not all labels. - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(1.0)), - )], - query, - 5_000, - WindowType::Tumbling, - ); - - let context = engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .expect("Failed to build context"); - - assert_eq!( - context.metadata.query_output_labels.labels, - vec!["host".to_string()], - "Collapsable query output labels should be spatial GROUP BY labels" - ); - } - - #[tokio::test] - async fn test_collapsable_empty_store() { - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![], - query, - 5_000, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert!( - results.is_empty(), - "Empty store should return 0 results, got {}", - results.len() - ); - } - - // ======================================================================== - // Non-collapsable spatial-of-temporal combinations must be rejected (#508) - // - // Before this fix, any spatial op could wrap any temporal function - // structurally, and a non-collapsable combination (e.g. `sum(min_over_time(...))`) - // silently dropped the outer aggregation instead of being rejected — returning - // ungrouped per-series values as if the outer `sum(...)` had never been - // written. `get_is_collapsable` only allows sum+sum_over_time, - // sum+count_over_time, min+min_over_time, max+max_over_time. - // ======================================================================== - - #[tokio::test] - async fn test_non_collapsable_sum_of_min_over_time_is_rejected() { - let query = "sum(min_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![], - query, - 5_000, - WindowType::Tumbling, - ); - - assert!( - engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .is_none(), - "sum+min_over_time is not collapsable and must not match any pattern" - ); - } - - #[tokio::test] - async fn test_non_collapsable_avg_of_rate_is_rejected() { - let query = "avg(rate(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![], - query, - 5_000, - WindowType::Tumbling, - ); - - assert!( - engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .is_none(), - "avg+rate is not collapsable and must not match any pattern" - ); - } - - // ======================================================================== - // Old-vs-New comparison tests for temporal queries - // ======================================================================== - - #[tokio::test] - async fn test_temporal_sum_over_time_old_vs_new() { - let data: TemporalData = TEMPORAL_TIMESTAMPS - .iter() - .map(|&ts| { - ( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - ) - }) - .collect(); - - // window_size_ms (1_000) must be smaller than the query's [5s] range so the - // derived do_merge (range_ms > window_size_ms) reflects a realistic config - // where the aggregation's bucket is finer than the requested range. - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 1_000, - WindowType::Tumbling, - ); - - assert_old_new_match(&engine, query, QUERY_TIME).await; - } - - #[tokio::test] - async fn test_collapsable_sum_of_sum_over_time_old_vs_new() { - let mut data: TemporalData = Vec::new(); - for &ts in &TEMPORAL_TIMESTAMPS { - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )); - data.push(( - ts, - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(20.0)) as Box, - )); - } - - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 1_000, - WindowType::Tumbling, - ); - - assert_old_new_match(&engine, query, QUERY_TIME).await; - } -} diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs deleted file mode 100644 index 8f942905..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs +++ /dev/null @@ -1,590 +0,0 @@ -//! Plan Execution Integration Tests -//! -//! Tests that verify the new DataFusion plan-based execution path -//! produces correct results for spatial queries. -//! -//! These tests use an actual store with test data. - -use crate::data_model::{AggregationType, KeyByLabelValues, Measurement}; -use crate::engines::simple_engine::SimpleEngine; -use crate::precompute_operators::sum_accumulator::SumAccumulator; -use std::collections::HashMap; - -#[cfg(test)] -mod tests { - use super::*; - use crate::precompute_operators::{ - DatasketchesKLLAccumulator, IncreaseAccumulator, MinMaxAccumulator, - MultipleMinMaxAccumulator, MultipleSumAccumulator, SetAggregatorAccumulator, - }; - use crate::tests::test_utilities::engine_factories::*; - - /// Creates a test engine and store with sample data for spatial sum queries - fn create_test_engine_with_data() -> SimpleEngine { - create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![ - ( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)), - ), - ( - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(200.0)), - ), - ], - "sum(http_requests) by (host)", - ) - } - - #[test] - fn test_to_logical_plan_produces_valid_structure() { - let engine = create_test_engine_with_data(); - - let query_time_sec = 1000.0; - let context = engine - .build_query_execution_context_promql( - "sum(http_requests) by (host)".to_string(), - query_time_sec, - ) - .expect("Failed to build context"); - - let plan = context.to_logical_plan(); - assert!( - plan.is_ok(), - "Failed to build logical plan: {:?}", - plan.err() - ); - - let plan = plan.unwrap(); - match &plan { - datafusion::logical_expr::LogicalPlan::Extension(ext) => { - assert_eq!( - ext.node.name(), - "SummaryInfer", - "Root should be SummaryInfer" - ); - } - _ => panic!("Expected Extension node at root, got {:?}", plan), - } - } - - #[tokio::test] - async fn test_execute_plan_returns_results() { - let engine = create_test_engine_with_data(); - - let query_time_sec = 1000.0; - let context = engine - .build_query_execution_context_promql( - "sum(http_requests) by (host)".to_string(), - query_time_sec, - ) - .expect("Failed to build context"); - - let result = engine.execute_plan(&context).await; - assert!(result.is_ok(), "execute_plan failed: {:?}", result.err()); - - let results = result.unwrap(); - assert_eq!( - results.len(), - 2, - "Expected 2 results, got {}", - results.len() - ); - - let values: Vec = results.iter().map(|r| r.value).collect(); - assert!( - values.contains(&100.0) && values.contains(&200.0), - "Expected values [100.0, 200.0], got {:?}", - values - ); - } - - #[tokio::test] - async fn test_execute_plan_correct_labels() { - let engine = create_test_engine_with_data(); - - let query_time_sec = 1000.0; - let context = engine - .build_query_execution_context_promql( - "sum(http_requests) by (host)".to_string(), - query_time_sec, - ) - .expect("Failed to build context"); - - let results = engine - .execute_plan(&context) - .await - .expect("execute_plan failed"); - - let result_map: HashMap = results - .iter() - .map(|r| (r.labels.labels[0].clone(), r.value)) - .collect(); - - assert_eq!(result_map.get("host-a"), Some(&100.0)); - assert_eq!(result_map.get("host-b"), Some(&200.0)); - } - - #[tokio::test] - async fn test_execute_plan_multiple_timestamps() { - let engine = create_engine_multi_timestamp( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![ - ( - 999_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(50.0)), - ), - ( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(50.0)), - ), - ], - "sum(http_requests) by (host)", - ); - - let context = engine - .build_query_execution_context_promql( - "sum(http_requests) by (host)".to_string(), - 1000.0, - ) - .expect("Failed to build context"); - - let results = engine - .execute_plan(&context) - .await - .expect("execute_plan failed"); - - assert_eq!(results.len(), 1, "Expected 1 result"); - assert_eq!(results[0].labels.labels[0], "host-a"); - assert_eq!( - results[0].value, 50.0, - "Spatial-only queries use latest timestamp only" - ); - } - - #[tokio::test] - async fn test_execute_plan_matches_old_pipeline() { - let engine = create_test_engine_with_data(); - assert_old_new_match(&engine, "sum(http_requests) by (host)", 1000.0).await; - } - - // ======================================================================== - // Category 1 - Old-vs-New Comparison for more accumulator types - // ======================================================================== - - #[tokio::test] - async fn test_old_vs_new_kll_quantile() { - let mut kll_a = DatasketchesKLLAccumulator::new(200); - for v in [10.0, 20.0, 30.0, 40.0, 50.0] { - kll_a.update(v); - } - let mut kll_b = DatasketchesKLLAccumulator::new(200); - for v in [100.0, 200.0, 300.0] { - kll_b.update(v); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(kll_a)), - (Some(vec!["host-b".to_string()]), Box::new(kll_b)), - ], - "quantile(0.5, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(0.5, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_kll_quantile_p99() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in 1..=100 { - kll.update(v as f64); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], - "quantile(0.99, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(0.99, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_kll_quantile_p0() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in [5.0, 10.0, 15.0, 20.0, 25.0] { - kll.update(v); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], - "quantile(0.0, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(0.0, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_kll_quantile_p1() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in [5.0, 10.0, 15.0, 20.0, 25.0] { - kll.update(v); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], - "quantile(1.0, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(1.0, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_kll_quantile_p25() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in 1..=1000 { - kll.update(v as f64); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], - "quantile(0.25, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(0.25, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - #[ignore = "Blocked: SetAggregatorAccumulator does not support old pipeline query"] - async fn test_old_vs_new_set_aggregator_cardinality() { - let mut set_a = SetAggregatorAccumulator::new(); - set_a.add_key(KeyByLabelValues { - labels: vec!["user-1".to_string()], - }); - set_a.add_key(KeyByLabelValues { - labels: vec!["user-2".to_string()], - }); - let mut set_b = SetAggregatorAccumulator::new(); - set_b.add_key(KeyByLabelValues { - labels: vec!["user-3".to_string()], - }); - - let engine = create_engine_single_pop( - "active_users", - AggregationType::SetAggregator, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(set_a)), - (Some(vec!["host-b".to_string()]), Box::new(set_b)), - ], - "count(active_users) by (host)", - ); - - assert_old_new_match(&engine, "count(active_users) by (host)", 1000.0).await; - } - - #[tokio::test] - #[ignore = "Blocked: IncreaseAccumulator has no arroyo serde"] - async fn test_old_vs_new_increase() { - let inc_a = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(100.0), 10); - let inc_b = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(200.0), 10); - - let engine = create_engine_single_pop( - "http_requests_total", - AggregationType::Increase, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(inc_a)), - (Some(vec!["host-b".to_string()]), Box::new(inc_b)), - ], - "sum(increase(http_requests_total[10s])) by (host)", - ); - - assert_old_new_match( - &engine, - "sum(increase(http_requests_total[10s])) by (host)", - 1000.0, - ) - .await; - } - - #[tokio::test] - #[ignore = "Blocked: MinMaxAccumulator has no arroyo serde"] - async fn test_old_vs_new_minmax_min() { - let mm_a = MinMaxAccumulator::with_value(10.0, "min".to_string()).unwrap(); - let mm_b = MinMaxAccumulator::with_value(5.0, "min".to_string()).unwrap(); - - let engine = create_engine_single_pop( - "temperature", - AggregationType::MinMax, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(mm_a)), - (Some(vec!["host-b".to_string()]), Box::new(mm_b)), - ], - "min(temperature) by (host)", - ); - - assert_old_new_match(&engine, "min(temperature) by (host)", 1000.0).await; - } - - #[tokio::test] - #[ignore = "Blocked: MinMaxAccumulator has no arroyo serde"] - async fn test_old_vs_new_minmax_max() { - let mm_a = MinMaxAccumulator::with_value(90.0, "max".to_string()).unwrap(); - let mm_b = MinMaxAccumulator::with_value(95.0, "max".to_string()).unwrap(); - - let engine = create_engine_single_pop( - "temperature", - AggregationType::MinMax, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(mm_a)), - (Some(vec!["host-b".to_string()]), Box::new(mm_b)), - ], - "max(temperature) by (host)", - ); - - assert_old_new_match(&engine, "max(temperature) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_multiple_sum() { - let mut ms = MultipleSumAccumulator::new(); - let key = KeyByLabelValues { - labels: vec!["host-a".to_string()], - }; - ms.update(key, 42.0); - - let engine = create_engine_single_pop( - "requests", - AggregationType::MultipleSum, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(ms))], - "sum(requests) by (host)", - ); - - assert_old_new_match(&engine, "sum(requests) by (host)", 1000.0).await; - } - - #[tokio::test] - #[ignore = "Blocked: MultipleMinMaxAccumulator has no arroyo serde"] - async fn test_old_vs_new_multiple_minmax() { - let mm = MultipleMinMaxAccumulator::new("min".to_string()).unwrap(); - - let engine = create_engine_single_pop( - "latency", - AggregationType::MultipleMinMax, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(mm))], - "min(latency) by (host)", - ); - - assert_old_new_match(&engine, "min(latency) by (host)", 1000.0).await; - } - - // ======================================================================== - // Category 3 - Edge Cases - // ======================================================================== - - #[tokio::test] - async fn test_execute_plan_empty_store() { - let engine = create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![], // No data - "sum(http_requests) by (host)", - ); - - let results = execute_new_plan(&engine, "sum(http_requests) by (host)", 1000.0).await; - assert!( - results.is_empty(), - "Empty store should return 0 results, got {}", - results.len() - ); - } - - #[tokio::test] - async fn test_execute_plan_many_groups() { - #[allow(clippy::type_complexity)] - let data: Vec<(Option>, Box)> = (0..100) - .map(|i| { - ( - Some(vec![format!("host-{:03}", i)]), - Box::new(SumAccumulator::with_sum(i as f64)) as Box, - ) - }) - .collect(); - - let engine = create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - "sum(http_requests) by (host)", - ); - - let results = execute_new_plan(&engine, "sum(http_requests) by (host)", 1000.0).await; - assert_eq!( - results.len(), - 100, - "Expected 100 results, got {}", - results.len() - ); - } - - #[tokio::test] - async fn test_execute_plan_multi_timestamp_multi_group() { - // 3 timestamps x 3 groups; spatial-only uses latest timestamp only -> 3 results - let mut data = Vec::new(); - for ts in [998_000u64, 999_000, 1_000_000] { - for i in 0..3 { - data.push(( - ts, - Some(vec![format!("host-{}", i)]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )); - } - } - - let engine = create_engine_multi_timestamp( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - "sum(http_requests) by (host)", - ); - - let results = execute_new_plan(&engine, "sum(http_requests) by (host)", 1000.0).await; - assert_eq!( - results.len(), - 3, - "3 groups merged across 3 timestamps should give 3 results" - ); - // Spatial-only queries use latest timestamp only, so each group = 10.0 - for r in &results { - assert!( - (r.value - 10.0).abs() < 1e-10, - "Each group should be 10.0 (latest timestamp only), got {}", - r.value - ); - } - } - - #[tokio::test] - async fn test_execute_plan_single_group_many_timestamps() { - // Spatial-only queries only consider the latest timestamp (query time). - // Include data at the query time (1_000_000) plus older timestamps. - #[allow(clippy::type_complexity)] - let mut data: Vec<(u64, Option>, Box)> = (0..9) - .map(|i| { - ( - (991_000 + i * 1000) as u64, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(5.0)) as Box, - ) - }) - .collect(); - data.push(( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(5.0)) as Box, - )); - - let engine = create_engine_multi_timestamp( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - "sum(http_requests) by (host)", - ); - - let results = execute_new_plan(&engine, "sum(http_requests) by (host)", 1000.0).await; - assert_eq!(results.len(), 1, "Single group should give 1 result"); - assert!( - (results[0].value - 5.0).abs() < 1e-10, - "Spatial-only uses latest timestamp only, expected 5.0, got {}", - results[0].value - ); - } - - // ======================================================================== - // Category 8 - Error Paths - // ======================================================================== - - #[tokio::test] - async fn test_execute_plan_not_implemented_increase() { - // IncreaseAccumulator has no arroyo serde, so execute_plan should fail. - // Uses a bare OnlyTemporal query (not `sum(increase(...))`): Increase is - // never collapsable with any spatial aggregation (see get_is_collapsable), - // so a spatial wrapper would no longer match any pattern at all (#508). - let inc = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(100.0), 10); - - let engine = create_engine_single_pop( - "http_requests_total", - AggregationType::Increase, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(inc))], - "increase(http_requests_total[10s])", - ); - - let context = engine - .build_query_execution_context_promql( - "increase(http_requests_total[10s])".to_string(), - 1000.0, - ) - .expect("Should build context"); - - let result = engine.execute_plan(&context).await; - assert!( - result.is_err(), - "execute_plan should fail for IncreaseAccumulator (no arroyo serde)" - ); - } - - #[tokio::test] - async fn test_execute_plan_not_implemented_minmax() { - let mm = MinMaxAccumulator::with_value(42.0, "min".to_string()).unwrap(); - - let engine = create_engine_single_pop( - "temperature", - AggregationType::MinMax, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(mm))], - "min(temperature) by (host)", - ); - - let context = engine - .build_query_execution_context_promql("min(temperature) by (host)".to_string(), 1000.0) - .expect("Should build context"); - - let result = engine.execute_plan(&context).await; - assert!( - result.is_err(), - "execute_plan should fail for MinMaxAccumulator (no arroyo serde)" - ); - } -} diff --git a/asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs b/asap-query-engine/src/tests/dispatch_arithmetic_tests.rs similarity index 100% rename from asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs rename to asap-query-engine/src/tests/dispatch_arithmetic_tests.rs diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 8fbd71ec..dcf04274 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -1,6 +1,6 @@ pub mod capability_matching_tests; pub mod clickhouse_forwarding_tests; -pub mod datafusion; +pub mod dispatch_arithmetic_tests; pub mod elastic_dsl_query_tests; pub mod elastic_forwarding_tests; pub mod native_binary_arithmetic_plan_tests; @@ -9,8 +9,10 @@ pub mod native_pipeline_merge_tests; pub mod native_range_query_tests; pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; +pub mod range_query_arithmetic_tests; pub mod sql_pattern_matching_tests; pub mod store_correctness_tests; +pub mod structural_matching_tests; pub mod trait_design_tests; #[cfg(test)] diff --git a/asap-query-engine/src/tests/native_range_query_tests.rs b/asap-query-engine/src/tests/native_range_query_tests.rs index 3f72ad99..0e838154 100644 --- a/asap-query-engine/src/tests/native_range_query_tests.rs +++ b/asap-query-engine/src/tests/native_range_query_tests.rs @@ -108,7 +108,7 @@ mod tests { type TimeSeriesData = Vec<(u64, Option>, Box)>; /// Dual-population counterpart to - /// datafusion::range_query_arithmetic_tests::create_range_engine_two_metrics: + /// range_query_arithmetic_tests::create_range_engine_two_metrics: /// one metric, separate value/key aggregations, data spread across /// multiple 1s tumbling-window buckets so a range query has more than /// one output step to expand keys for. diff --git a/asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs b/asap-query-engine/src/tests/range_query_arithmetic_tests.rs similarity index 100% rename from asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs rename to asap-query-engine/src/tests/range_query_arithmetic_tests.rs diff --git a/asap-query-engine/src/tests/datafusion/structural_matching_tests.rs b/asap-query-engine/src/tests/structural_matching_tests.rs similarity index 100% rename from asap-query-engine/src/tests/datafusion/structural_matching_tests.rs rename to asap-query-engine/src/tests/structural_matching_tests.rs 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 ab44a9b0..b3c54f5c 100644 --- a/asap-query-engine/src/tests/test_utilities/engine_factories.rs +++ b/asap-query-engine/src/tests/test_utilities/engine_factories.rs @@ -10,7 +10,6 @@ use crate::data_model::{ KeyByLabelValues, PrecomputedOutput, PromQLSchema, QueryConfig, QueryLanguage, SchemaConfig, StreamingConfig, WindowType, }; -use crate::engines::query_result::InstantVectorElement; use crate::engines::simple_engine::SimpleEngine; use crate::stores::simple_map_store::SimpleMapStore; use crate::stores::Store; @@ -630,74 +629,3 @@ pub fn create_engine_multi_timestamp_with_window( QueryLanguage::promql, ) } - -/// Execute both old pipeline and new plan-based path, compare results with epsilon tolerance. -pub async fn assert_old_new_match(engine: &SimpleEngine, query: &str, query_time_sec: f64) { - let context = engine - .build_query_execution_context_promql(query.to_string(), query_time_sec) - .expect("Failed to build context"); - - let old_results = engine - .execute_query_pipeline(&context, false, false) - .expect("Old pipeline failed"); - - let new_results = engine - .execute_plan(&context) - .await - .expect("New plan path failed"); - - assert_eq!( - old_results.len(), - new_results.len(), - "Result count mismatch: old={}, new={}", - old_results.len(), - new_results.len() - ); - - let old_map: HashMap, f64> = old_results - .iter() - .map(|r| (r.labels.labels.clone(), r.value)) - .collect(); - - let new_map: HashMap, f64> = new_results - .iter() - .map(|r| (r.labels.labels.clone(), r.value)) - .collect(); - - for (key, old_value) in &old_map { - let new_value = new_map - .get(key) - .unwrap_or_else(|| panic!("Key {:?} missing from new results", key)); - assert!( - (old_value - new_value).abs() < 1e-10, - "Value mismatch for key {:?}: old={}, new={}", - key, - old_value, - new_value - ); - } - - for key in new_map.keys() { - assert!( - old_map.contains_key(key), - "Extra key {:?} in new results", - key - ); - } -} - -/// Convenience wrapper to execute via the new plan path. -pub async fn execute_new_plan( - engine: &SimpleEngine, - query: &str, - query_time_sec: f64, -) -> Vec { - let context = engine - .build_query_execution_context_promql(query.to_string(), query_time_sec) - .expect("Failed to build context"); - - engine - .execute_plan(&context) - .await - .expect("execute_plan failed") -}