From c1fceba4a5c071934c8db1bcafaf1c64515b2a32 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Thu, 2 Jul 2026 11:55:42 -0400 Subject: [PATCH 1/2] fix(query-engine): added topk detection support for SpatioTemporal queries --- .../src/engines/simple_engine/sql.rs | 140 +++++++++++++++++- 1 file changed, 137 insertions(+), 3 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/sql.rs b/asap-query-engine/src/engines/simple_engine/sql.rs index ba2c70b..2947bbc 100644 --- a/asap-query-engine/src/engines/simple_engine/sql.rs +++ b/asap-query-engine/src/engines/simple_engine/sql.rs @@ -699,15 +699,30 @@ impl SimpleEngine { .get_name() .to_lowercase(); - let statistic_to_compute = Self::parse_single_statistic(&statistic_name)?; + // SpatioTemporal queries are a single (non-nested) SELECT layer, same + // shape `detect_sql_topk` expects, so top-k detection applies directly. + let topk = detect_sql_topk(query_data); + if topk.is_some_and(|t| t.weighting == TopkWeighting::Sum) { + warn!( + "SUM top-k assumes non-negative values; results are undefined for columns with negative entries" + ); + } + let statistic_to_compute = if topk.is_some() { + Statistic::Topk + } else { + Self::parse_single_statistic(&statistic_name)? + }; - let query_kwargs = self + let mut query_kwargs = self .build_query_kwargs_sql(&statistic_to_compute, match_result) .map_err(|e| { warn!("{}", e); e }) .ok()?; + if let Some(topk) = topk { + query_kwargs.insert("k".to_string(), topk.k.to_string()); + } let metadata = QueryMetadata { query_output_labels: query_output_labels.clone(), @@ -743,7 +758,7 @@ impl SimpleEngine { let requirements = self.build_query_requirements_sql( match_result, QueryPatternType::OnlyTemporal, - None, + topk, ); self.streaming_config .read() @@ -1276,6 +1291,96 @@ mod topk_pipeline_tests { ) } + /// Build a SQL engine whose only aggregation is a self-keyed + /// `CountMinSketchWithHeap` precomputed over the *full 2-second* window, + /// referenced by a query_config template with a matching 2s duration — + /// same shape as `build_topk_engine`, just sized to a SpatioTemporal + /// window instead of a single-scrape-interval one. + /// + /// Self-keyed resolution (`aggregation_id_for_key == + /// aggregation_id_for_value`) only happens via this query_config path + /// (`find_query_config_sql` / `get_aggregation_id_info`'s single-reference + /// case). The capability-matching fallback always pairs + /// `CountMinSketchWithHeap` with a separate key aggregation (see + /// `count_topk_capability_fallback_pairs_heap_with_key_agg`) — it doesn't + /// know a heap can be self-keyed, so a SpatioTemporal top-k query with no + /// matching query_config would fail to resolve today. That gap is + /// tracked separately; this test targets `build_spatiotemporal_context`'s + /// top-k *detection* (issue #498), not the capability-matching fallback. + fn build_spatiotemporal_topk_engine() -> (SimpleEngine, Arc) { + let template = "SELECT srcip, COUNT(pkt_len) FROM netflow_table \ + WHERE time BETWEEN DATEADD(s, -2, NOW()) AND NOW() GROUP BY srcip"; + + let value_cols: HashSet = ["pkt_len"].iter().map(|s| s.to_string()).collect(); + let labels: HashSet = ["srcip", "dstip", "proto"] + .iter() + .map(|s| s.to_string()) + .collect(); + let table = Table::new(METRIC.to_string(), "time".to_string(), value_cols, labels); + let sql_schema = SQLSchema::new(vec![table]); + + let query_config = QueryConfig::new(template.to_string()) + .add_aggregation(AggregationReference::new(AGG_ID, None)); + + let inference_config = InferenceConfig { + schema: SchemaConfig::SQL(sql_schema), + query_configs: vec![query_config], + cleanup_policy: CleanupPolicy::NoCleanup, + }; + + let agg_config = AggregationConfig { + aggregation_id: AGG_ID, + aggregation_type: AggregationType::CountMinSketchWithHeap, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::empty(), + aggregated_labels: KeyByLabelNames::new(vec!["srcip".to_string()]), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: 2000, + slide_interval_ms: 2000, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: METRIC.to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }; + + let mut agg_configs = HashMap::new(); + agg_configs.insert(AGG_ID, agg_config); + let streaming_config = Arc::new(StreamingConfig { + aggregation_configs: agg_configs, + }); + + let store = Arc::new(SimpleMapStore::new( + streaming_config.clone(), + CleanupPolicy::NoCleanup, + )); + + let engine = SimpleEngine::new( + store.clone(), + inference_config, + streaming_config, + 1000, // 1s scrape interval ⇒ a 2s query window classifies as SpatioTemporal + QueryLanguage::sql, + ); + (engine, store) + } + + /// Incoming top-k query over a 2-second window grouped by a *subset* of + /// labels (`srcip` only, out of `srcip`/`dstip`/`proto`) — the shape that + /// classifies as `SpatioTemporal` rather than `OnlySpatial`. + fn spatiotemporal_topk_query(limit: u64) -> String { + format!( + "SELECT srcip, COUNT(pkt_len) AS transfer_events FROM netflow_table \ + WHERE time BETWEEN DATEADD(s, -2, '2025-10-01 00:00:10') AND '2025-10-01 00:00:10' \ + GROUP BY srcip ORDER BY transfer_events DESC LIMIT {limit}" + ) + } + /// Incoming SUM top-k query over a 1-second absolute window. fn sum_topk_query(limit: u64) -> String { format!( @@ -1506,6 +1611,35 @@ mod topk_pipeline_tests { assert!(context.store_plan.keys_query.is_none()); } + /// `build_spatiotemporal_context` (issue #498) must run the same top-k + /// detection as the `OnlyTemporal`/`OnlySpatial` path: a `SpatioTemporal` + /// query (multi-interval window, subset of labels) shaped like + /// `COUNT ... GROUP BY ORDER BY DESC LIMIT k` still resolves + /// to `Statistic::Topk` with `k` threaded through, self-keyed to the same + /// sketch the query_config template resolves for plain COUNT. + #[test] + fn spatiotemporal_query_detects_topk_and_resolves_self_keyed_heap() { + let (engine, _store) = build_spatiotemporal_topk_engine(); + let context = engine + .build_query_execution_context_sql(spatiotemporal_topk_query(10), QUERY_TIME) + .expect("SpatioTemporal top-k query should build a context via the query_config path"); + + assert_eq!( + context.metadata.statistic_to_compute, + Statistic::Topk, + "ORDER BY DESC LIMIT n must be promoted to Topk even under SpatioTemporal classification", + ); + assert_eq!( + context.metadata.query_kwargs.get("k").map(String::as_str), + Some("10"), + "LIMIT should be threaded through as the `k` kwarg", + ); + assert_eq!( + context.agg_info.aggregation_id_for_key, context.agg_info.aggregation_id_for_value, + "self-keyed: the heap supplies both keys and counts", + ); + } + #[test] fn returns_top_k_srcips_sorted_descending() { let (engine, store) = build_topk_engine(); From a90edb5153f0a11c3ab44e3be7729b0f5947a3d9 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Thu, 2 Jul 2026 12:54:19 -0400 Subject: [PATCH 2/2] added another test Co-Authored-By: Claude Sonnet 5 --- asap-planner-rs/tests/sql_integration.rs | 38 ++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/asap-planner-rs/tests/sql_integration.rs b/asap-planner-rs/tests/sql_integration.rs index 8b7bb1b..465f144 100644 --- a/asap-planner-rs/tests/sql_integration.rs +++ b/asap-planner-rs/tests/sql_integration.rs @@ -978,6 +978,44 @@ fn spatial_sum_topk_heap() { ); } +/// COUNT … ORDER BY DESC LIMIT k over a *multi-scrape-interval* window +/// (2s duration, 1s ingest ⇒ `QueryType::SpatioTemporal`, not `Spatial`) still +/// detects as top-k. `detect_sql_topk` (called unconditionally on any +/// non-nested query, see `planner/sql.rs`) isn't gated by `QueryType`, unlike +/// the query engine's `build_spatiotemporal_context` (issue #498) — this test +/// locks in that the planner already covers the SpatioTemporal case. +#[test] +fn spatiotemporal_count_topk_heap() { + let q = "SELECT srcip, COUNT(pkt_len) AS transfer_events FROM netflow_table \ + WHERE time BETWEEN DATEADD(s, -12, NOW()) AND DATEADD(s, -10, NOW()) \ + GROUP BY srcip ORDER BY transfer_events DESC LIMIT 10"; + let out = SQLController::from_yaml(&netflow_one_query_config(q, 2_000), sql_opts_1s_ingest()) + .unwrap() + .generate() + .unwrap(); + + assert_eq!(out.streaming_aggregation_count(), 1); + assert_eq!(out.inference_query_count(), 1); + assert!(out.has_aggregation_type("CountMinSketchWithHeap")); + assert!(out.has_aggregation_type_and_sub_type("CountMinSketchWithHeap", "topk")); + assert!(!out.has_aggregation_type("DeltaSetAggregator")); + assert!(!out.has_aggregation_type("CountMinSketch")); + assert!(out.all_tumbling_window_sizes_eq(2_000)); + assert_eq!( + out.aggregation_labels("CountMinSketchWithHeap", "grouping"), + Vec::::new() + ); + assert_eq!( + out.aggregation_labels("CountMinSketchWithHeap", "aggregated"), + vec!["srcip".to_string()] + ); + assert_eq!( + out.aggregation_parameter("CountMinSketchWithHeap", "count_events") + .and_then(|v| v.as_bool()), + Some(true) + ); +} + /// Plain COUNT without ORDER BY / LIMIT stays on the CMS + DeltaSet path. #[test] fn spatial_count_without_order_by_is_not_topk() {