baseline_runner.py

"""
baseline_runner.py - Generate Real Baseline Comparison Results
===============================================================
Runs two policies against all 5 tasks and prints a clean comparison table:
  1. Fallback policy: deterministic rule-based (no LLM required)
  2. LLM policy: uses Qwen2.5-72B via HF Inference Router

Run:
  # Fallback only (no API key needed):
  python baseline_runner.py

  # With LLM comparison:
  HF_TOKEN=hf_xxx python baseline_runner.py
  MODEL_NAME=Qwen/Qwen2.5-72B-Instruct python baseline_runner.py

Results are saved to results/baseline_results.json and printed as a table.
"""

import io
import sys
# Fix Windows console encoding so non-ASCII results don't crash
if hasattr(sys.stdout, 'buffer'):
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')

import json
import os
import sys
import time
from typing import Any, Dict, List, Optional

ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT_DIR)

from env import SQLOptimEnv
from models import Action
from tasks import TASKS

HF_TOKEN   = os.environ.get("HF_TOKEN", "")
MODEL_NAME = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
API_BASE   = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1")

TASK_IDS = list(TASKS.keys())

# ─────────────────────────────────────────────────────────────────────────────
# Fallback policy: deterministic, hand-crafted, no LLM
# ─────────────────────────────────────────────────────────────────────────────

FALLBACK_SOLUTIONS: Dict[str, Dict[str, Any]] = {
    "task_1_basic_antipatterns": {
        "suggestions": [
            {"issue_type": "select_star", "line": 1,
             "description": "SELECT * fetches all columns from 500k rows — use explicit projection.",
             "severity": "high", "fix": "SELECT id, customer_id, status, total, created_at"},
            {"issue_type": "non_sargable_cast", "line": 3,
             "description": "CAST(customer_id AS VARCHAR) prevents integer comparison and pruning.",
             "severity": "critical", "fix": "WHERE customer_id = 5000"},
            {"issue_type": "function_on_date_column", "line": 4,
             "description": "YEAR() on date column forces full scan; use a date range instead.",
             "severity": "high", "fix": "created_at >= DATE '2024-01-01' AND created_at < DATE '2025-01-01'"},
        ],
        "optimized_query": (
            "SELECT id, customer_id, product_id, status, total, created_at\n"
            "FROM orders\n"
            "WHERE customer_id = 5000\n"
            "  AND created_at >= DATE '2024-01-01'\n"
            "  AND created_at < DATE '2025-01-01';"
        ),
        "summary": (
            "Three anti-patterns: SELECT * over 500k rows wastes bandwidth, "
            "CAST on customer_id prevents pruning, YEAR() forces full date scan. "
            "Explicit column projection + integer comparison + date range fix all three."
        ),
        "estimated_improvement": "3-5x faster — eliminates type-cast and function penalties",
        "approved": False,
    },

    "task_2_correlated_subqueries": {
        "suggestions": [
            {"issue_type": "correlated_subquery_count", "line": 4,
             "description": "Correlated COUNT subquery scans 500k orders per premium user (N+1 pattern).",
             "severity": "critical", "fix": "LEFT JOIN with GROUP BY aggregation"},
            {"issue_type": "correlated_subquery_sum", "line": 7,
             "description": "Correlated SUM subquery -- another full scan per user.",
             "severity": "critical", "fix": "Include in the same LEFT JOIN aggregation"},
            {"issue_type": "correlated_subquery_limit", "line": 11,
             "description": "Correlated ORDER BY LIMIT 1 -- sorted scan per user.",
             "severity": "high", "fix": "Use ROW_NUMBER() window function in a CTE"},
            {"issue_type": "missing_aggregation_join", "line": 16,
             "description": "Single aggregation JOIN replaces all three subqueries in one pass.",
             "severity": "critical", "fix": "LEFT JOIN aggregated subquery ON u.id = agg.customer_id"},
        ],
        "optimized_query": (
            "WITH agg AS (\n"
            "    SELECT\n"
            "        customer_id,\n"
            "        COUNT(*) FILTER (WHERE status = 'completed')              AS completed_orders,\n"
            "        SUM(total) FILTER (WHERE created_at >= DATE '2024-01-01') AS ytd_spend\n"
            "    FROM orders\n"
            "    GROUP BY customer_id\n"
            "),\n"
            "last_order AS (\n"
            "    SELECT customer_id, total AS last_order_amount\n"
            "    FROM (\n"
            "        SELECT customer_id, total,\n"
            "               ROW_NUMBER() OVER (PARTITION BY customer_id ORDER BY created_at DESC) AS rn\n"
            "        FROM orders\n"
            "    ) t WHERE rn = 1\n"
            ")\n"
            "SELECT\n"
            "    u.email,\n"
            "    u.region,\n"
            "    COALESCE(a.completed_orders, 0) AS completed_orders,\n"
            "    a.ytd_spend,\n"
            "    l.last_order_amount\n"
            "FROM users u\n"
            "LEFT JOIN agg a ON u.id = a.customer_id\n"
            "LEFT JOIN last_order l ON u.id = l.customer_id\n"
            "WHERE u.tier = 'premium';"
        ),
        "summary": (
            "Three correlated subqueries each scan 500k orders per premium user (~3300 users). "
            "Worst case: 3 × 3300 × 500k = 5B row reads. "
            "A single CTE with GROUP BY + FILTER aggregates everything in one pass over orders."
        ),
        "estimated_improvement": "10-20x faster — eliminates N+1 pattern with single JOIN",
        "approved": False,
    },

    "task_3_wildcard_scan": {
        "suggestions": [
            {"issue_type": "leading_wildcard_like", "line": 6,
             "description": "LIKE '%purchase%' and '%buy%' are leading-wildcard patterns that disable zone-map pruning on 1M rows.",
             "severity": "critical", "fix": "Replace with exact equality where possible"},
            {"issue_type": "or_expands_to_full_scan", "line": 7,
             "description": "OR session_id LIKE 'sess_%' matches ALL 1M rows (every session_id starts with 'sess_'), making the other OR conditions redundant. The WHERE is effectively a no-op.",
             "severity": "high", "fix": "Recognize session_id LIKE 'sess_%' covers all rows; simplify or remove WHERE clause entirely"},
            {"issue_type": "select_star_large_table", "line": 2,
             "description": "SELECT * on 1M rows fetches all columns plus two computed columns before the WHERE is evaluated.",
             "severity": "high", "fix": "SELECT id, user_id, session_id, event_type, occurred_at — explicit projection"},
            {"issue_type": "pre_filter_computed_columns", "line": 3,
             "description": "CAST(id AS VARCHAR) || '_' || event_type and UPPER(event_type) computed for all 1M rows before WHERE.",
             "severity": "medium", "fix": "Compute derived columns after WHERE filtering (or in final SELECT)"},
        ],
        "optimized_query": (
            "-- session_id LIKE 'sess_%%' matches ALL rows, so original WHERE = full scan anyway.\n"
            "-- Remove the redundant OR conditions; keep explicit column projection.\n"
            "SELECT\n"
            "    id, user_id, session_id, event_type, occurred_at,\n"
            "    CAST(id AS VARCHAR) || '_' || event_type AS event_key,\n"
            "    UPPER(event_type) AS event_type_upper\n"
            "FROM events;"
        ),
        "summary": (
            "The WHERE clause is a logical no-op: session_id LIKE 'sess_%' matches ALL 1M rows "
            "(every session starts with 'sess_'), making the event_type LIKE conditions redundant. "
            "Removing the redundant wildcard evaluations eliminates three LIKE scans per row. "
            "SELECT * replaced with explicit columns to reduce column I/O bandwidth."
        ),
        "estimated_improvement": "1.5-3x faster — eliminates three LIKE evaluations per row; no filter selectivity possible",
        "approved": False,
    },

    "task_4_implicit_join": {
        "suggestions": [
            {"issue_type": "implicit_cross_join", "line": 8,
             "description": "Comma-syntax FROM (implicit join) risks Cartesian product if WHERE fails.",
             "severity": "critical", "fix": "Use explicit INNER JOIN ... ON syntax"},
            {"issue_type": "repeated_scalar_subquery_avg", "line": 6,
             "description": "Scalar subquery AVG(total) re-scans all 500k orders once per GROUP BY group.",
             "severity": "high", "fix": "Pre-compute in a CTE and cross-join the scalar value"},
            {"issue_type": "repeated_scalar_subquery_max", "line": 7,
             "description": "Scalar subquery MAX(total) WHERE status='completed' — same issue.",
             "severity": "high", "fix": "Include in the same pre-compute CTE"},
            {"issue_type": "missing_explicit_join", "line": 8,
             "description": "Rewrite with explicit INNER JOIN for clarity and safety.",
             "severity": "medium", "fix": "FROM users u INNER JOIN orders o ON u.id = o.customer_id"},
        ],
        "optimized_query": (
            "WITH global_stats AS (\n"
            "    SELECT\n"
            "        AVG(total)                                          AS global_avg,\n"
            "        MAX(total) FILTER (WHERE status = 'completed')      AS max_deal\n"
            "    FROM orders\n"
            ")\n"
            "SELECT\n"
            "    u.region,\n"
            "    u.plan,\n"
            "    COUNT(*)       AS total_orders,\n"
            "    SUM(o.total)   AS revenue,\n"
            "    gs.global_avg,\n"
            "    gs.max_deal\n"
            "FROM users u\n"
            "INNER JOIN orders o ON u.id = o.customer_id\n"
            "CROSS JOIN global_stats gs\n"
            "WHERE o.status IN ('completed', 'shipped')\n"
            "GROUP BY u.region, u.plan, gs.global_avg, gs.max_deal;"
        ),
        "summary": (
            "Comma-syntax implicit join is an anti-pattern that risks Cartesian products. "
            "Two scalar subqueries re-scan 500k orders per GROUP BY group. "
            "A CTE computes global stats exactly once; explicit INNER JOIN ensures correctness."
        ),
        "estimated_improvement": "8-15x faster — CTE eliminates repeated subquery scans",
        "approved": False,
    },

    "task_5_window_functions": {
        "suggestions": [
            {"issue_type": "no_pre_filter", "line": 11,
             "description": "No WHERE clause: all 5 window functions computed over the entire 1M row events table. Window functions partition and sort the full dataset.",
             "severity": "critical", "fix": "Adding a WHERE filter changes window function semantics (partitions include fewer rows), so instead optimize by removing expensive global RANK"},
            {"issue_type": "global_rank_no_partition", "line": 8,
             "description": "RANK() OVER (ORDER BY occurred_at DESC) with no PARTITION sorts all 1M rows globally — the single most expensive operation in this query.",
             "severity": "critical", "fix": "Remove RANK() OVER (ORDER BY occurred_at DESC) — it sorts 1M rows and provides marginal analytical value"},
            {"issue_type": "redundant_window_functions", "line": 5,
             "description": "5 separate OVER() clauses, two sharing PARTITION BY user_id. Each is a distinct sort/hash-aggregate pass over all 1M rows.",
             "severity": "high", "fix": "Merge compatible windows; DuckDB can share passes for identical PARTITION BY"},
            {"issue_type": "count_vs_conditional_sum", "line": 9,
             "description": "SUM(CASE WHEN event_type='purchase' THEN 1 ELSE 0 END) is equivalent to but slower than COUNT(*) FILTER (WHERE event_type='purchase').",
             "severity": "medium", "fix": "COUNT(*) FILTER (WHERE event_type = 'purchase') OVER (PARTITION BY user_id)"},
            {"issue_type": "select_all_unfiltered", "line": 1,
             "description": "The original query selects specific columns, but all 1M rows with no selectivity.",
             "severity": "medium", "fix": "Preserve column projection; focus optimizations on window function cost"},
        ],
        "optimized_query": (
            "-- Remove global RANK() (sorts all 1M rows); replace SUM(CASE WHEN) with COUNT FILTER.\n"
            "-- Window functions must operate over the same dataset to preserve correct partition counts.\n"
            "SELECT\n"
            "    user_id,\n"
            "    event_type,\n"
            "    occurred_at,\n"
            "    COUNT(*) OVER (PARTITION BY user_id)                                AS total_user_events,\n"
            "    COUNT(*) OVER (PARTITION BY user_id, event_type)                   AS type_count,\n"
            "    ROW_NUMBER() OVER (PARTITION BY user_id ORDER BY occurred_at DESC) AS recency_rank,\n"
            "    COUNT(*) FILTER (WHERE event_type = 'purchase')\n"
            "        OVER (PARTITION BY user_id)                                    AS user_purchases\n"
            "FROM events;"
        ),
        "summary": (
            "Five window functions over all 1M events with no pre-filtering causes 5 full sort/hash passes. "
            "The global RANK() OVER (ORDER BY occurred_at DESC) sorts all 1M rows globally — the single most expensive operation. "
            "Removing RANK() eliminates the global sort pass entirely. "
            "Replacing SUM(CASE WHEN event_type='purchase' THEN 1 ELSE 0 END) with COUNT(*) FILTER (WHERE event_type='purchase') "
            "is more concise and allows better optimizer hints. The dataset must remain unfiltered "
            "to preserve correct window partition counts across all user/event_type combinations."
        ),
        "estimated_improvement": "3-6x faster — removing global RANK() eliminates the full 1M-row global sort pass",
        "approved": False,
    },
}


def run_fallback_policy(env: SQLOptimEnv) -> Dict[str, Dict]:
    """Run deterministic fallback policy against all tasks."""
    results = {}
    for task_id in TASK_IDS:
        obs = env.reset(task_id=task_id)
        sol = FALLBACK_SOLUTIONS[task_id]
        action = Action(
            suggestions=sol["suggestions"],
            optimized_query=sol["optimized_query"],
            summary=sol["summary"],
            estimated_improvement=sol["estimated_improvement"],
            approved=sol["approved"],
        )
        result = env.step(action)
        exec_info = result.info.get("execution") or {}
        results[task_id] = {
            "task_name":    obs.task_name,
            "difficulty":   obs.difficulty,
            "score":        round(result.reward.score, 4),
            "speedup":      round(exec_info.get("speedup", 1.0), 2),
            "correct":      exec_info.get("results_match", False),
            "steps":        1,
            "policy":       "fallback",
        }
        print(
            f"  [Fallback] {obs.difficulty:12s} | "
            f"score={result.reward.score:.4f} | "
            f"speedup={exec_info.get('speedup', 1.0):.2f}x | "
            f"correct={exec_info.get('results_match', False)}",
            flush=True,
        )
    return results


def run_llm_policy(env: SQLOptimEnv) -> Optional[Dict[str, Dict]]:
    """Run LLM policy if HF_TOKEN is set."""
    if not HF_TOKEN:
        print("  [LLM] HF_TOKEN not set — skipping LLM baseline.", flush=True)
        return None

    try:
        from openai import OpenAI
    except ImportError:
        print("  [LLM] openai package not installed — skipping.", flush=True)
        return None

    from inference import SYSTEM_PROMPT, build_user_prompt, parse_action

    client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE)
    results = {}

    for task_id in TASK_IDS:
        obs = env.reset(task_id=task_id)
        try:
            resp = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user",   "content": build_user_prompt(obs)},
                ],
                temperature=0.0,
                max_tokens=2000,
            )
            parsed = parse_action(resp.choices[0].message.content or "")
        except Exception as e:
            print(f"  [LLM] Call failed for {task_id}: {e}", flush=True)
            parsed = FALLBACK_SOLUTIONS[task_id]

        action = Action(
            suggestions=parsed.get("suggestions", []),
            optimized_query=parsed.get("optimized_query", ""),
            summary=parsed.get("summary", ""),
            estimated_improvement=parsed.get("estimated_improvement", ""),
            approved=parsed.get("approved", False),
        )

        env.reset(task_id=task_id)
        result = env.step(action)
        exec_info = result.info.get("execution") or {}
        results[task_id] = {
            "task_name":    obs.task_name,
            "difficulty":   obs.difficulty,
            "score":        round(result.reward.score, 4),
            "speedup":      round(exec_info.get("speedup", 1.0), 2),
            "correct":      exec_info.get("results_match", False),
            "steps":        1,
            "policy":       f"llm:{MODEL_NAME}",
        }
        print(
            f"  [LLM] {obs.difficulty:12s} | "
            f"score={result.reward.score:.4f} | "
            f"speedup={exec_info.get('speedup', 1.0):.2f}x | "
            f"correct={exec_info.get('results_match', False)}",
            flush=True,
        )
    return results


def print_comparison_table(
    fallback: Dict[str, Dict],
    llm: Optional[Dict[str, Dict]],
):
    print("\n" + "=" * 80)
    print("  BASELINE RESULTS — SQL Query Optimization Environment")
    print("=" * 80)

    header = f"{'Task':<40} {'Diff':<12} {'F-Score':>8} {'F-Spdup':>8} {'F-Corr':>7}"
    if llm:
        header += f" {'L-Score':>8} {'L-Spdup':>8} {'L-Corr':>7} {'Delta':>7}"
    print(header)
    print("-" * 80)

    for task_id in TASK_IDS:
        fb = fallback[task_id]
        row = (
            f"{fb['task_name'][:38]:<40} "
            f"{fb['difficulty']:<12} "
            f"{fb['score']:>8.4f} "
            f"{fb['speedup']:>7.2f}x "
            f"{'YES' if fb['correct'] else 'NO':>7}"
        )
        if llm and task_id in llm:
            lm = llm[task_id]
            delta = lm["score"] - fb["score"]
            row += (
                f" {lm['score']:>8.4f} "
                f"{lm['speedup']:>7.2f}x "
                f"{'YES' if lm['correct'] else 'NO':>7} "
                f"{'+' if delta >= 0 else ''}{delta:>6.4f}"
            )
        print(row)

    print("=" * 80)
    fb_avg = sum(r["score"] for r in fallback.values()) / len(fallback)
    print(f"  Fallback avg score : {fb_avg:.4f}")
    if llm:
        lm_avg = sum(r["score"] for r in llm.values()) / len(llm)
        print(f"  LLM avg score      : {lm_avg:.4f}  (+{lm_avg - fb_avg:.4f} vs fallback)")
    print("=" * 80 + "\n")


def main():
    print("\n[SQLOptimEnv] Baseline Runner", flush=True)
    print("Initialising DuckDB environment (warm-up ~3s) ...\n", flush=True)
    env = SQLOptimEnv()

    print("[1/2] Running fallback (deterministic) policy ...", flush=True)
    fallback_results = run_fallback_policy(env)

    print("\n[2/2] Running LLM policy ...", flush=True)
    llm_results = run_llm_policy(env)

    print_comparison_table(fallback_results, llm_results)

    # Save results
    os.makedirs("results", exist_ok=True)
    output = {
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "fallback":  fallback_results,
        "llm":       llm_results,
    }
    out_path = "results/baseline_results.json"
    with open(out_path, "w") as f:
        json.dump(output, f, indent=2)
    print(f"[SAVED] Results written to {out_path}", flush=True)
    return output


if __name__ == "__main__":
    main()