fix(asap-tools): score SQL query results in calculate_fidelity

asap-common/dependencies/py/promql_utilities/promql_utilities/query_results/classes.py

@@ -1,6 +1,84 @@
+import re
+
 import numpy as np
 from typing import List, Dict, Optional, Set
 
+# Parens, commas and the keywords that bound a SELECT list.
+_SELECT_LIST_TOKEN = re.compile(r"\(|\)|,|\bSELECT\b|\bFROM\b", re.IGNORECASE)
+_IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
+
+
+def _outer_select_items(query: str) -> List[str]:
+    """Items of the outermost SELECT list.
+
+    e.g. "SELECT job, QUANTILE(0.9, value) FROM t GROUP BY job"
+    -> ["job", "QUANTILE(0.9, value)"]
+    """
+    depth = 0
+    item_start: Optional[int] = None
+    items: List[str] = []
+    for token in _SELECT_LIST_TOKEN.finditer(query):
+        word = token.group().upper()
+        if word == "(":
+            depth += 1
+        elif word == ")":
+            depth -= 1
+        elif depth > 0:
+            continue
+        elif item_start is None:
+            if word == "SELECT":
+                item_start = token.end()
+        elif word in (",", "FROM"):
+            items.append(query[item_start : token.start()].strip())
+            if word == "FROM":
+                return items
+            item_start = token.end()
+    raise ValueError(f"No top-level SELECT ... FROM in query: {query!r}")
+
+
+def _sql_column_label_names(query: str) -> List[Optional[str]]:
+    """Label name for each result column, None for the single value column.
+
+    Bare column names in the SELECT list are the GROUP BY labels; the one
+    other item is the aggregate. Names come from the query, not the result,
+    because TabSeparated output has no header.
+    """
+    names = [
+        item if _IDENTIFIER.fullmatch(item) else None
+        for item in _outer_select_items(query)
+    ]
+    if names.count(None) != 1:
+        raise ValueError(
+            f"Expected exactly one non-label column in SELECT list of {query!r}"
+        )
+    return names
+
+
+def parse_tabseparated_result(query: str, raw_text: str) -> Dict[frozenset, float]:
+    """Key ClickHouse TabSeparated rows like `QueryResult.result`.
+
+    e.g. for "SELECT job, SUM(value) FROM t GROUP BY job", the row "api\t3"
+    becomes {frozenset({("job", "api")}): 3.0}. Columns are read in SELECT
+    list order, as ClickHouse returns them.
+    """
+    names = _sql_column_label_names(query)
+    value_idx = names.index(None)
+    result: Dict[frozenset, float] = {}
+    for row in raw_text.splitlines():
+        fields = row.split("\t")
+        if len(fields) != len(names):
+            raise ValueError(
+                f"Row {row!r} has {len(fields)} columns, query {query!r} "
+                f"selects {len(names)}"
+            )
+        key = frozenset(
+            (name, field) for name, field in zip(names, fields) if name is not None
+        )
+        if key in result:
+            raise ValueError(f"Duplicate label set {sorted(key)} in SQL result")
+        result[key] = float(fields[value_idx])
+    return result
+
 
 class TimeSeries:
     def __init__(self, key: frozenset, values: List[Optional[float]]):
@@ -39,6 +117,12 @@ def __init__(
                 for result_per_key in result
             }
 
+    def keyed_result(self) -> Optional[Dict[frozenset, float]]:
+        """`result`, or for SQL queries the parsed `raw_text_result`."""
+        if self.raw_text_result is None:
+            return self.result
+        return parse_tabseparated_result(self.query, self.raw_text_result)
+
 
 class QueryResultAcrossTime:
     def __init__(self, server_name, query, query_idx, num_repetitions):
@@ -52,10 +136,11 @@ def add_result(self, query_result: QueryResult):
         self.query_results.append(query_result)
 
     def get_all_timeseries(self) -> Dict[frozenset, TimeSeries]:
+        keyed_results = [qr.keyed_result() for qr in self.query_results]
         keys: Set[frozenset] = set()
-        for query_result in self.query_results:
-            if query_result.result:
-                keys.update(query_result.result.keys())
+        for keyed_result in keyed_results:
+            if keyed_result:
+                keys.update(keyed_result.keys())
 
         assert len(self.query_results) == self.num_repetitions
         ret: Dict[frozenset, TimeSeries] = {}
@@ -64,11 +149,11 @@ def get_all_timeseries(self) -> Dict[frozenset, TimeSeries]:
         }
 
         for k in keys:
-            for repetition_idx, result in enumerate(self.query_results):
+            for repetition_idx, keyed_result in enumerate(keyed_results):
                 # A key absent from one repetition (e.g. a series that appears
                 # mid-run) stays None for that repetition.
-                if result.result:
-                    intermediate_ret[k][repetition_idx] = result.result.get(k)
+                if keyed_result:
+                    intermediate_ret[k][repetition_idx] = keyed_result.get(k)
 
             ret[k] = TimeSeries(k, intermediate_ret[k])
 

asap-tools/experiments/tests/test_calculate_fidelity.py

@@ -0,0 +1,170 @@
+"""Tests for scoring SQL (ClickHouse TabSeparated) query results."""
+
+import math
+import os
+import sys
+import unittest
+
+sys.path.insert(
+    0,
+    os.path.join(
+        os.path.dirname(os.path.abspath(__file__)),
+        "..",
+        "..",
+        "..",
+        "asap-common",
+        "dependencies",
+        "py",
+        "promql_utilities",
+    ),
+)
+
+from promql_utilities.query_results.classes import (  # noqa: E402
+    QueryResult,
+    QueryResultAcrossTime,
+    parse_tabseparated_result,
+)
+from post_experiment.single_experiment.calculate_fidelity import (  # noqa: E402
+    get_timeseries_similarity_scores,
+    mape,
+)
+
+GROUPED_QUERY = (
+    "SELECT job, instance, QUANTILE(0.9, value) FROM data "
+    "WHERE time BETWEEN DATEADD(s, -300, NOW()) AND NOW() GROUP BY job, instance"
+)
+
+
+class ParseTabSeparatedResultTest(unittest.TestCase):
+    def test_group_by_columns_become_labels_named_from_select_list(self):
+        parsed = parse_tabseparated_result(
+            GROUPED_QUERY, "api\ta:1\t3.5\nweb\tb:2\t7\n"
+        )
+        self.assertEqual(
+            parsed,
+            {
+                frozenset({("job", "api"), ("instance", "a:1")}): 3.5,
+                frozenset({("job", "web"), ("instance", "b:2")}): 7.0,
+            },
+        )
+
+    def test_value_column_is_located_by_select_position(self):
+        query = "SELECT SUM(value) AS v, job FROM data GROUP BY job"
+        self.assertEqual(
+            parse_tabseparated_result(query, "4\tapi\n"),
+            {frozenset({("job", "api")}): 4.0},
+        )
+
+    def test_commas_inside_nested_query_do_not_split_outer_select(self):
+        query = (
+            "SELECT job, MAX(s) FROM (SELECT job, instance, SUM(value) AS s "
+            "FROM data GROUP BY job, instance) GROUP BY job"
+        )
+        self.assertEqual(
+            parse_tabseparated_result(query, "api\t9\n"),
+            {frozenset({("job", "api")}): 9.0},
+        )
+
+    def test_single_value_without_group_by(self):
+        query = "SELECT QUANTILE(0.5, value) FROM data"
+        self.assertEqual(
+            parse_tabseparated_result(query, "42.5\n"), {frozenset(): 42.5}
+        )
+
+    def test_empty_result(self):
+        self.assertEqual(parse_tabseparated_result(GROUPED_QUERY, ""), {})
+
+    def test_nan_and_inf_values(self):
+        # ClickHouse prints nan/inf/-inf; ASAPQuery's adapter prints Rust's NaN.
+        parsed = parse_tabseparated_result(
+            "SELECT job, AVG(value) FROM data GROUP BY job",
+            "a\tnan\nb\tinf\nc\t-inf\nd\tNaN\n",
+        )
+        self.assertTrue(math.isnan(parsed[frozenset({("job", "a")})]))
+        self.assertEqual(parsed[frozenset({("job", "b")})], math.inf)
+        self.assertEqual(parsed[frozenset({("job", "c")})], -math.inf)
+        self.assertTrue(math.isnan(parsed[frozenset({("job", "d")})]))
+
+    def test_non_numeric_value_fails(self):
+        with self.assertRaises(ValueError):
+            parse_tabseparated_result(GROUPED_QUERY, "api\ta:1\tnot-a-number\n")
+
+    def test_column_count_mismatch_fails(self):
+        with self.assertRaises(ValueError):
+            parse_tabseparated_result(GROUPED_QUERY, "api\t3.5\n")
+
+    def test_duplicate_label_set_fails(self):
+        with self.assertRaises(ValueError):
+            parse_tabseparated_result(GROUPED_QUERY, "api\ta:1\t1\napi\ta:1\t2\n")
+
+    def test_select_list_without_exactly_one_aggregate_fails(self):
+        with self.assertRaises(ValueError):
+            parse_tabseparated_result(
+                "SELECT job, SUM(value), MAX(value) FROM data GROUP BY job", ""
+            )
+
+    def test_query_without_select_from_fails(self):
+        with self.assertRaises(ValueError):
+            parse_tabseparated_result("rate(x[1m])", "")
+
+
+def _sql_results(server_name, raw_texts):
+    across_time = QueryResultAcrossTime(server_name, GROUPED_QUERY, 0, len(raw_texts))
+    for repetition_idx, raw_text in enumerate(raw_texts):
+        across_time.add_result(
+            QueryResult(
+                server_name,
+                GROUPED_QUERY,
+                0,
+                repetition_idx,
+                result=None,
+                latency=0.1,
+                cumulative_latency=None,
+                raw_text_result=raw_text,
+            )
+        )
+    return across_time
+
+
+class SqlFidelityTest(unittest.TestCase):
+    def test_sql_query_is_scored(self):
+        exact = _sql_results("clickhouse", ["api\ta:1\t10\n", "api\ta:1\t20\n"])
+        estimate = _sql_results("asap", ["api\ta:1\t11\n", "api\ta:1\t22\n"])
+
+        scores = get_timeseries_similarity_scores(
+            [exact], [estimate], [GROUPED_QUERY], [mape], verbose=False
+        )
+
+        self.assertAlmostEqual(scores["mape"][GROUPED_QUERY], 10.0)
+
+    def test_sql_exact_matches_prometheus_json_estimate(self):
+        # An estimate answered as Prometheus JSON keys series by label name,
+        # so it must line up with the SELECT-list names given to the TSV.
+        exact = _sql_results("clickhouse", ["api\ta:1\t10\n"])
+        estimate = QueryResultAcrossTime("asap", GROUPED_QUERY, 0, 1)
+        estimate.add_result(
+            QueryResult(
+                "asap",
+                GROUPED_QUERY,
+                0,
+                0,
+                result=[
+                    {
+                        "metric": {"instance": "a:1", "job": "api"},
+                        "value": [0, "12"],
+                    }
+                ],
+                latency=0.1,
+                cumulative_latency=None,
+            )
+        )
+
+        scores = get_timeseries_similarity_scores(
+            [exact], [estimate], [GROUPED_QUERY], [mape], verbose=False
+        )
+
+        self.assertAlmostEqual(scores["mape"][GROUPED_QUERY], 20.0)
+
+
+if __name__ == "__main__":
+    unittest.main()