scripts/metrics.py

"""Metrics for the Squirrel-mini toy reproduction.

Implements:
  - script complexity stats (lines, tokens, function count, AST width/depth) for
    comparison with Table 1 of the paper.
  - Exact Match (EM) and a Graph-Match PROXY (GM-proxy) for comparison with
    Table 2. The paper's GM uses Apache Calcite optimized-plan graph isomorphism;
    we approximate it with sqlglot's parsed+optimized AST compared structurally,
    which is a weaker but self-contained proxy (documented as a substitution).
"""
from __future__ import annotations

import re
from typing import Optional

import sqlglot
import sqlglot.expressions as sge
import tiktoken

_ENC = tiktoken.get_encoding("cl100k_base")

FUNC_RE = re.compile(r"\b([A-Za-z_][A-Za-z0-9_]*)\s*\(")
SQL_KEYWORDS_NOT_FUNCS = {
    "select", "from", "where", "group", "order", "having", "join", "on", "as",
    "case", "when", "then", "partition", "over", "union", "insert", "into",
    "values", "with",
}


def count_lines(sql: str) -> int:
    return len([ln for ln in sql.splitlines() if ln.strip()])


def count_tokens(sql: str) -> int:
    return len(_ENC.encode(sql))


def count_functions(sql: str) -> int:
    matches = FUNC_RE.findall(sql)
    return sum(1 for m in matches if m.lower() not in SQL_KEYWORDS_NOT_FUNCS)


def ast_width_depth(sql: str, dialect: str = "hive") -> tuple[Optional[int], Optional[int]]:
    """Return (max width, max depth) of the sqlglot AST tree, as a proxy for the
    paper's Apache-Calcite-based AST width/depth (Eq. 2)."""
    try:
        tree = sqlglot.parse_one(sql, read=dialect)
    except Exception:
        return None, None

    depths: dict[int, int] = {}
    max_depth = 0

    def children(node):
        for value in node.args.values():
            items = value if isinstance(value, list) else [value]
            for item in items:
                if isinstance(item, sge.Expression):
                    yield item

    def walk(node, depth=0):
        nonlocal max_depth
        max_depth = max(max_depth, depth)
        depths[depth] = depths.get(depth, 0) + 1
        for child in children(node):
            walk(child, depth + 1)

    walk(tree, 0)
    max_width = max(depths.values()) if depths else 0
    return max_width, max_depth


def script_stats(sql: str, dialect: str = "hive") -> dict:
    width, depth = ast_width_depth(sql, dialect)
    return {
        "lines": count_lines(sql),
        "tokens": count_tokens(sql),
        "functions": count_functions(sql),
        "ast_width": width,
        "ast_depth": depth,
    }


def normalize_sql(sql: str) -> str:
    return re.sub(r"\s+", " ", sql.strip().lower())


def exact_match(pred: str, ref: str) -> bool:
    return normalize_sql(pred) == normalize_sql(ref)


_COMMENT_RE = re.compile(r"--[^\n]*|/\*.*?\*/", re.S)


def strip_comments(sql: str) -> str:
    return _COMMENT_RE.sub(" ", sql)


def graph_match_proxy(pred: str, ref: str, dialect: str = "hive") -> bool:
    """Weak proxy for the paper's Graph Match score: parse both queries with
    sqlglot, canonicalize (strip comments, normalize identifiers), and compare
    the resulting AST dumps for structural equality. Falls back to False if
    either side fails to parse.

    Comments are stripped before comparison because the paper's real GM score
    operates on Apache Calcite's optimized relational-algebra plan, which has
    no notion of SQL comments at all (unlike its strict, string-level EM
    score) -- this is exactly the EM-vs-GM gap the paper's Table 2 is built
    around, so a comment-only diff should read as a GM hit, not a miss."""
    try:
        pred_tree = sqlglot.parse_one(strip_comments(pred), read=dialect)
        ref_tree = sqlglot.parse_one(strip_comments(ref), read=dialect)
    except Exception:
        return False
    try:
        pred_norm = pred_tree.sql(dialect=dialect, normalize=True, pretty=False)
        ref_norm = ref_tree.sql(dialect=dialect, normalize=True, pretty=False)
    except Exception:
        pred_norm, ref_norm = str(pred_tree), str(ref_tree)
    if normalize_sql(pred_norm) == normalize_sql(ref_norm):
        return True
    # Fall back to comparing the raw AST repr (catches semantically-irrelevant
    # whitespace/casing differences that .sql() normalization might not fully collapse)
    return pred_tree.dump() == ref_tree.dump()


def modify_better_proxy(pred: str, ref: str, buggy: str, dialect: str = "hive") -> Optional[bool]:
    """Weak proxy for Modify-Better score: character-edit-distance based since a
    full AST edit-distance is out of scope for this toy harness."""
    import difflib

    def dist(a: str, b: str) -> float:
        return 1 - difflib.SequenceMatcher(None, a, b).ratio()

    ref_n, pred_n, buggy_n = normalize_sql(ref), normalize_sql(pred), normalize_sql(buggy)
    return dist(pred_n, ref_n) < dist(buggy_n, ref_n)