Skip to content

Commit 3ff12e2

Browse files
fix(asap-tools): score SQL query results in calculate_fidelity (#896)
SQL results arrive as ClickHouse TabSeparated text in QueryResult.raw_text_result with `result` left None, so get_all_timeseries found no series and SQL queries were never scored. Parse the TSV into the same {frozenset(label -> value): float} form as `result`. TSV has no header, so label names come from the query's outermost SELECT list (bare columns are labels, the single other item is the value), read in SELECT order as ClickHouse returns them. Parsing happens in get_all_timeseries via QueryResult.keyed_result(), so both freshly-run and JSONL-loaded results are covered. Claude-Session: https://claude.ai/code/session_011gKTCAEz3sJ1Hzde16xxU1 Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 81fb7c4 commit 3ff12e2

2 files changed

Lines changed: 261 additions & 6 deletions

File tree

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

Lines changed: 91 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,84 @@
1+
import re
2+
13
import numpy as np
24
from typing import List, Dict, Optional, Set
35

6+
# Parens, commas and the keywords that bound a SELECT list.
7+
_SELECT_LIST_TOKEN = re.compile(r"\(|\)|,|\bSELECT\b|\bFROM\b", re.IGNORECASE)
8+
_IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
9+
10+
11+
def _outer_select_items(query: str) -> List[str]:
12+
"""Items of the outermost SELECT list.
13+
14+
e.g. "SELECT job, QUANTILE(0.9, value) FROM t GROUP BY job"
15+
-> ["job", "QUANTILE(0.9, value)"]
16+
"""
17+
depth = 0
18+
item_start: Optional[int] = None
19+
items: List[str] = []
20+
for token in _SELECT_LIST_TOKEN.finditer(query):
21+
word = token.group().upper()
22+
if word == "(":
23+
depth += 1
24+
elif word == ")":
25+
depth -= 1
26+
elif depth > 0:
27+
continue
28+
elif item_start is None:
29+
if word == "SELECT":
30+
item_start = token.end()
31+
elif word in (",", "FROM"):
32+
items.append(query[item_start : token.start()].strip())
33+
if word == "FROM":
34+
return items
35+
item_start = token.end()
36+
raise ValueError(f"No top-level SELECT ... FROM in query: {query!r}")
37+
38+
39+
def _sql_column_label_names(query: str) -> List[Optional[str]]:
40+
"""Label name for each result column, None for the single value column.
41+
42+
Bare column names in the SELECT list are the GROUP BY labels; the one
43+
other item is the aggregate. Names come from the query, not the result,
44+
because TabSeparated output has no header.
45+
"""
46+
names = [
47+
item if _IDENTIFIER.fullmatch(item) else None
48+
for item in _outer_select_items(query)
49+
]
50+
if names.count(None) != 1:
51+
raise ValueError(
52+
f"Expected exactly one non-label column in SELECT list of {query!r}"
53+
)
54+
return names
55+
56+
57+
def parse_tabseparated_result(query: str, raw_text: str) -> Dict[frozenset, float]:
58+
"""Key ClickHouse TabSeparated rows like `QueryResult.result`.
59+
60+
e.g. for "SELECT job, SUM(value) FROM t GROUP BY job", the row "api\t3"
61+
becomes {frozenset({("job", "api")}): 3.0}. Columns are read in SELECT
62+
list order, as ClickHouse returns them.
63+
"""
64+
names = _sql_column_label_names(query)
65+
value_idx = names.index(None)
66+
result: Dict[frozenset, float] = {}
67+
for row in raw_text.splitlines():
68+
fields = row.split("\t")
69+
if len(fields) != len(names):
70+
raise ValueError(
71+
f"Row {row!r} has {len(fields)} columns, query {query!r} "
72+
f"selects {len(names)}"
73+
)
74+
key = frozenset(
75+
(name, field) for name, field in zip(names, fields) if name is not None
76+
)
77+
if key in result:
78+
raise ValueError(f"Duplicate label set {sorted(key)} in SQL result")
79+
result[key] = float(fields[value_idx])
80+
return result
81+
482

583
class TimeSeries:
684
def __init__(self, key: frozenset, values: List[Optional[float]]):
@@ -39,6 +117,12 @@ def __init__(
39117
for result_per_key in result
40118
}
41119

120+
def keyed_result(self) -> Optional[Dict[frozenset, float]]:
121+
"""`result`, or for SQL queries the parsed `raw_text_result`."""
122+
if self.raw_text_result is None:
123+
return self.result
124+
return parse_tabseparated_result(self.query, self.raw_text_result)
125+
42126

43127
class QueryResultAcrossTime:
44128
def __init__(self, server_name, query, query_idx, num_repetitions):
@@ -52,10 +136,11 @@ def add_result(self, query_result: QueryResult):
52136
self.query_results.append(query_result)
53137

54138
def get_all_timeseries(self) -> Dict[frozenset, TimeSeries]:
139+
keyed_results = [qr.keyed_result() for qr in self.query_results]
55140
keys: Set[frozenset] = set()
56-
for query_result in self.query_results:
57-
if query_result.result:
58-
keys.update(query_result.result.keys())
141+
for keyed_result in keyed_results:
142+
if keyed_result:
143+
keys.update(keyed_result.keys())
59144

60145
assert len(self.query_results) == self.num_repetitions
61146
ret: Dict[frozenset, TimeSeries] = {}
@@ -64,11 +149,11 @@ def get_all_timeseries(self) -> Dict[frozenset, TimeSeries]:
64149
}
65150

66151
for k in keys:
67-
for repetition_idx, result in enumerate(self.query_results):
152+
for repetition_idx, keyed_result in enumerate(keyed_results):
68153
# A key absent from one repetition (e.g. a series that appears
69154
# mid-run) stays None for that repetition.
70-
if result.result:
71-
intermediate_ret[k][repetition_idx] = result.result.get(k)
155+
if keyed_result:
156+
intermediate_ret[k][repetition_idx] = keyed_result.get(k)
72157

73158
ret[k] = TimeSeries(k, intermediate_ret[k])
74159

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
"""Tests for scoring SQL (ClickHouse TabSeparated) query results."""
2+
3+
import math
4+
import os
5+
import sys
6+
import unittest
7+
8+
sys.path.insert(
9+
0,
10+
os.path.join(
11+
os.path.dirname(os.path.abspath(__file__)),
12+
"..",
13+
"..",
14+
"..",
15+
"asap-common",
16+
"dependencies",
17+
"py",
18+
"promql_utilities",
19+
),
20+
)
21+
22+
from promql_utilities.query_results.classes import ( # noqa: E402
23+
QueryResult,
24+
QueryResultAcrossTime,
25+
parse_tabseparated_result,
26+
)
27+
from post_experiment.single_experiment.calculate_fidelity import ( # noqa: E402
28+
get_timeseries_similarity_scores,
29+
mape,
30+
)
31+
32+
GROUPED_QUERY = (
33+
"SELECT job, instance, QUANTILE(0.9, value) FROM data "
34+
"WHERE time BETWEEN DATEADD(s, -300, NOW()) AND NOW() GROUP BY job, instance"
35+
)
36+
37+
38+
class ParseTabSeparatedResultTest(unittest.TestCase):
39+
def test_group_by_columns_become_labels_named_from_select_list(self):
40+
parsed = parse_tabseparated_result(
41+
GROUPED_QUERY, "api\ta:1\t3.5\nweb\tb:2\t7\n"
42+
)
43+
self.assertEqual(
44+
parsed,
45+
{
46+
frozenset({("job", "api"), ("instance", "a:1")}): 3.5,
47+
frozenset({("job", "web"), ("instance", "b:2")}): 7.0,
48+
},
49+
)
50+
51+
def test_value_column_is_located_by_select_position(self):
52+
query = "SELECT SUM(value) AS v, job FROM data GROUP BY job"
53+
self.assertEqual(
54+
parse_tabseparated_result(query, "4\tapi\n"),
55+
{frozenset({("job", "api")}): 4.0},
56+
)
57+
58+
def test_commas_inside_nested_query_do_not_split_outer_select(self):
59+
query = (
60+
"SELECT job, MAX(s) FROM (SELECT job, instance, SUM(value) AS s "
61+
"FROM data GROUP BY job, instance) GROUP BY job"
62+
)
63+
self.assertEqual(
64+
parse_tabseparated_result(query, "api\t9\n"),
65+
{frozenset({("job", "api")}): 9.0},
66+
)
67+
68+
def test_single_value_without_group_by(self):
69+
query = "SELECT QUANTILE(0.5, value) FROM data"
70+
self.assertEqual(
71+
parse_tabseparated_result(query, "42.5\n"), {frozenset(): 42.5}
72+
)
73+
74+
def test_empty_result(self):
75+
self.assertEqual(parse_tabseparated_result(GROUPED_QUERY, ""), {})
76+
77+
def test_nan_and_inf_values(self):
78+
# ClickHouse prints nan/inf/-inf; ASAPQuery's adapter prints Rust's NaN.
79+
parsed = parse_tabseparated_result(
80+
"SELECT job, AVG(value) FROM data GROUP BY job",
81+
"a\tnan\nb\tinf\nc\t-inf\nd\tNaN\n",
82+
)
83+
self.assertTrue(math.isnan(parsed[frozenset({("job", "a")})]))
84+
self.assertEqual(parsed[frozenset({("job", "b")})], math.inf)
85+
self.assertEqual(parsed[frozenset({("job", "c")})], -math.inf)
86+
self.assertTrue(math.isnan(parsed[frozenset({("job", "d")})]))
87+
88+
def test_non_numeric_value_fails(self):
89+
with self.assertRaises(ValueError):
90+
parse_tabseparated_result(GROUPED_QUERY, "api\ta:1\tnot-a-number\n")
91+
92+
def test_column_count_mismatch_fails(self):
93+
with self.assertRaises(ValueError):
94+
parse_tabseparated_result(GROUPED_QUERY, "api\t3.5\n")
95+
96+
def test_duplicate_label_set_fails(self):
97+
with self.assertRaises(ValueError):
98+
parse_tabseparated_result(GROUPED_QUERY, "api\ta:1\t1\napi\ta:1\t2\n")
99+
100+
def test_select_list_without_exactly_one_aggregate_fails(self):
101+
with self.assertRaises(ValueError):
102+
parse_tabseparated_result(
103+
"SELECT job, SUM(value), MAX(value) FROM data GROUP BY job", ""
104+
)
105+
106+
def test_query_without_select_from_fails(self):
107+
with self.assertRaises(ValueError):
108+
parse_tabseparated_result("rate(x[1m])", "")
109+
110+
111+
def _sql_results(server_name, raw_texts):
112+
across_time = QueryResultAcrossTime(server_name, GROUPED_QUERY, 0, len(raw_texts))
113+
for repetition_idx, raw_text in enumerate(raw_texts):
114+
across_time.add_result(
115+
QueryResult(
116+
server_name,
117+
GROUPED_QUERY,
118+
0,
119+
repetition_idx,
120+
result=None,
121+
latency=0.1,
122+
cumulative_latency=None,
123+
raw_text_result=raw_text,
124+
)
125+
)
126+
return across_time
127+
128+
129+
class SqlFidelityTest(unittest.TestCase):
130+
def test_sql_query_is_scored(self):
131+
exact = _sql_results("clickhouse", ["api\ta:1\t10\n", "api\ta:1\t20\n"])
132+
estimate = _sql_results("asap", ["api\ta:1\t11\n", "api\ta:1\t22\n"])
133+
134+
scores = get_timeseries_similarity_scores(
135+
[exact], [estimate], [GROUPED_QUERY], [mape], verbose=False
136+
)
137+
138+
self.assertAlmostEqual(scores["mape"][GROUPED_QUERY], 10.0)
139+
140+
def test_sql_exact_matches_prometheus_json_estimate(self):
141+
# An estimate answered as Prometheus JSON keys series by label name,
142+
# so it must line up with the SELECT-list names given to the TSV.
143+
exact = _sql_results("clickhouse", ["api\ta:1\t10\n"])
144+
estimate = QueryResultAcrossTime("asap", GROUPED_QUERY, 0, 1)
145+
estimate.add_result(
146+
QueryResult(
147+
"asap",
148+
GROUPED_QUERY,
149+
0,
150+
0,
151+
result=[
152+
{
153+
"metric": {"instance": "a:1", "job": "api"},
154+
"value": [0, "12"],
155+
}
156+
],
157+
latency=0.1,
158+
cumulative_latency=None,
159+
)
160+
)
161+
162+
scores = get_timeseries_similarity_scores(
163+
[exact], [estimate], [GROUPED_QUERY], [mape], verbose=False
164+
)
165+
166+
self.assertAlmostEqual(scores["mape"][GROUPED_QUERY], 20.0)
167+
168+
169+
if __name__ == "__main__":
170+
unittest.main()

0 commit comments

Comments
 (0)