RonyForAI/Mirage_DB_RL
0
1"""2tasks.py — Enterprise-grade query optimization task definitions for Mirage_RL.3 4Each task tier (easy / medium / hard) contains a pool of production-realistic5query scenarios sampled randomly per episode so agents cannot memorize solutions.6 7Domains covered:8 - E-commerce (OLTP): orders, customers, products, inventory, suppliers9 - Analytics (OLAP): events, sessions, campaigns, conversions, attribution10 - Financial (OLAP): transactions, accounts, merchants, fraud signals, risk11 12Key design decisions:13 - Table cardinalities match realistic production scales (10K – 1B rows)14 - Cardinality estimation noise (log-normal) simulates the core real-world15 challenge: a planner's row estimates are always wrong. Noise sigma maps to:16 σ=0.04 excellent statistics (small, frequently vacuumed)17 σ=0.15 typical OLTP tables18 σ=0.25 typical analytics/OLAP tables19 σ=0.40 poor statistics (large tables, column correlations)20 σ=0.60 very poor (event tables, multi-column predicates, data skew)21"""22 23from __future__ import annotations24 25import math26import random27from dataclasses import dataclass, field28from typing import List, Optional29 30 31# ─────────────────────────────────────────────────────────────────────────────32# Schema building-blocks33# ─────────────────────────────────────────────────────────────────────────────34 35@dataclass36class TableSpec:37 """Configuration for one table in a query scenario."""38 name: str39 true_rows: int # ground-truth row count (used for actual cost)40 selectivity: float # join predicate selectivity (fraction of rows kept)41 has_index: int # 1 = covering/clustered index present, 0 = heap scan42 noise_sigma: float # log-normal σ for cardinality estimation error43 44 45@dataclass46class Scenario:47 """A complete multi-table query scenario."""48 name: str # short identifier49 domain: str # "ecommerce" | "analytics" | "financial" | "saas"50 query_context: str # SQL-style description shown to the agent51 tables: List[TableSpec]52 53 54@dataclass55class TaskConfig:56 task_id: str57 name: str58 difficulty: str # "easy" | "medium" | "hard"59 description: str60 scenarios: List[Scenario] # pool sampled randomly per episode61 62 63# ─────────────────────────────────────────────────────────────────────────────64# Production Scenario Pools65# ─────────────────────────────────────────────────────────────────────────────66 67# ── EASY (3 tables, clean statistics, all indexes present) ──────────────────68EASY_SCENARIOS: List[Scenario] = [69 Scenario(70 name="ecommerce_catalog_lookup",71 domain="ecommerce",72 query_context=(73 "SELECT p.name, c.label, s.region "74 "FROM products JOIN categories ON p.category_id=c.id "75 "JOIN suppliers ON p.supplier_id=s.id "76 "WHERE c.segment='Electronics' AND s.active=true"77 ),78 tables=[79 TableSpec("products", true_rows=500_000, selectivity=0.04, has_index=1, noise_sigma=0.05),80 TableSpec("categories", true_rows=10_000, selectivity=0.45, has_index=1, noise_sigma=0.03),81 TableSpec("suppliers", true_rows=50_000, selectivity=0.12, has_index=1, noise_sigma=0.04),82 ],83 ),84 Scenario(85 name="saas_active_subscriptions",86 domain="saas",87 query_context=(88 "SELECT u.email, a.plan, s.renewal_date "89 "FROM users JOIN accounts ON u.account_id=a.id "90 "JOIN subscriptions ON a.id=s.account_id "91 "WHERE a.tier='enterprise' AND s.status='active'"92 ),93 tables=[94 TableSpec("users", true_rows=2_000_000, selectivity=0.08, has_index=1, noise_sigma=0.05),95 TableSpec("accounts", true_rows=800_000, selectivity=0.15, has_index=1, noise_sigma=0.04),96 TableSpec("subscriptions", true_rows=1_200_000, selectivity=0.10, has_index=1, noise_sigma=0.04),97 ],98 ),99 Scenario(100 name="inventory_reorder_check",101 domain="ecommerce",102 query_context=(103 "SELECT p.sku, w.location, i.quantity "104 "FROM products JOIN warehouses ON i.warehouse_id=w.id "105 "JOIN inventory ON p.id=i.product_id "106 "WHERE i.quantity < p.reorder_point AND w.region='US-WEST'"107 ),108 tables=[109 TableSpec("products", true_rows=500_000, selectivity=0.05, has_index=1, noise_sigma=0.04),110 TableSpec("warehouses", true_rows=5_000, selectivity=0.55, has_index=1, noise_sigma=0.03),111 TableSpec("inventory", true_rows=2_500_000, selectivity=0.03, has_index=1, noise_sigma=0.06),112 ],113 ),114]115 116# ── MEDIUM (5 tables, realistic noise, mixed index coverage) ─────────────────117MEDIUM_SCENARIOS: List[Scenario] = [118 Scenario(119 name="ecommerce_order_fulfillment",120 domain="ecommerce",121 query_context=(122 "SELECT o.id, c.name, p.sku, cat.label, w.region "123 "FROM orders JOIN customers ON o.customer_id=c.id "124 "JOIN products ON o.product_id=p.id "125 "JOIN categories ON p.category_id=cat.id "126 "JOIN warehouses ON o.warehouse_id=w.id "127 "WHERE o.status='pending' AND o.created_at > NOW() - INTERVAL '7 days'"128 ),129 tables=[130 TableSpec("orders", true_rows=10_000_000, selectivity=0.08, has_index=1, noise_sigma=0.15),131 TableSpec("customers", true_rows=2_000_000, selectivity=0.12, has_index=1, noise_sigma=0.10),132 TableSpec("products", true_rows=500_000, selectivity=0.05, has_index=1, noise_sigma=0.08),133 TableSpec("categories", true_rows=10_000, selectivity=0.40, has_index=1, noise_sigma=0.05),134 TableSpec("warehouses", true_rows=5_000, selectivity=0.55, has_index=0, noise_sigma=0.05),135 ],136 ),137 Scenario(138 name="marketing_funnel_analytics",139 domain="analytics",140 query_context=(141 "SELECT s.id, u.segment, c.name, cv.revenue, ch.source "142 "FROM sessions JOIN users ON s.user_id=u.id "143 "JOIN campaigns ON s.campaign_id=c.id "144 "JOIN conversions ON s.id=cv.session_id "145 "JOIN channels ON c.channel_id=ch.id "146 "WHERE c.type='paid' AND cv.revenue > 0"147 ),148 tables=[149 TableSpec("sessions", true_rows=50_000_000, selectivity=0.05, has_index=1, noise_sigma=0.25),150 TableSpec("users", true_rows=5_000_000, selectivity=0.15, has_index=1, noise_sigma=0.12),151 TableSpec("campaigns", true_rows=100_000, selectivity=0.25, has_index=1, noise_sigma=0.07),152 TableSpec("conversions", true_rows=2_000_000, selectivity=0.08, has_index=1, noise_sigma=0.18),153 TableSpec("channels", true_rows=500, selectivity=0.70, has_index=1, noise_sigma=0.02),154 ],155 ),156 Scenario(157 name="financial_transaction_summary",158 domain="financial",159 query_context=(160 "SELECT t.amount, a.holder, m.name, rs.score, cur.symbol "161 "FROM transactions JOIN accounts ON t.account_id=a.id "162 "JOIN merchants ON t.merchant_id=m.id "163 "JOIN risk_scores ON a.id=rs.account_id "164 "JOIN currencies ON t.currency_code=cur.code "165 "WHERE t.created_at >= CURRENT_DATE - 30 AND a.status='active'"166 ),167 tables=[168 TableSpec("transactions", true_rows=100_000_000, selectivity=0.03, has_index=1, noise_sigma=0.20),169 TableSpec("accounts", true_rows=10_000_000, selectivity=0.10, has_index=1, noise_sigma=0.10),170 TableSpec("merchants", true_rows=500_000, selectivity=0.18, has_index=1, noise_sigma=0.08),171 TableSpec("risk_scores", true_rows=10_000_000, selectivity=0.08, has_index=0, noise_sigma=0.20),172 TableSpec("currencies", true_rows=200, selectivity=0.80, has_index=1, noise_sigma=0.02),173 ],174 ),175]176 177# ── HARD (7 tables, high noise, missing indexes, skewed large tables) ────────178HARD_SCENARIOS: List[Scenario] = [179 Scenario(180 name="ecommerce_full_pipeline_audit",181 domain="ecommerce",182 query_context=(183 "SELECT o.id, oi.qty, p.sku, cat.label, c.name, s.contact, w.region "184 "FROM orders JOIN order_items ON o.id=oi.order_id "185 "JOIN products ON oi.product_id=p.id "186 "JOIN categories ON p.category_id=cat.id "187 "JOIN customers ON o.customer_id=c.id "188 "JOIN suppliers ON p.supplier_id=s.id "189 "JOIN warehouses ON oi.warehouse_id=w.id "190 "WHERE o.created_at >= '2024-01-01' AND c.country='US' "191 "AND cat.segment='Electronics'"192 ),193 tables=[194 TableSpec("categories", true_rows=1_000_000, selectivity=0.10, has_index=1, noise_sigma=0.05), # 100K195 TableSpec("warehouses", true_rows=3_000_000, selectivity=0.10, has_index=0, noise_sigma=0.08), # 300K196 TableSpec("products", true_rows=10_000_000, selectivity=0.05, has_index=1, noise_sigma=0.10), # 500K197 TableSpec("suppliers", true_rows=8_000_000, selectivity=0.10, has_index=0, noise_sigma=0.12), # 800K198 TableSpec("customers", true_rows=10_000_000, selectivity=0.12, has_index=1, noise_sigma=0.15), # 1.2M199 TableSpec("orders", true_rows=20_000_000, selectivity=0.08, has_index=1, noise_sigma=0.20), # 1.6M200 TableSpec("order_items", true_rows=40_000_000, selectivity=0.05, has_index=0, noise_sigma=0.35), # 2.0M201 ],202 ),203 Scenario(204 name="fraud_detection_pipeline",205 domain="financial",206 query_context=(207 "SELECT t.id, a.holder, m.category, fl.label, rs.score, d.fingerprint, loc.country "208 "FROM transactions JOIN accounts ON t.account_id=a.id "209 "JOIN merchants ON t.merchant_id=m.id "210 "JOIN fraud_labels ON t.id=fl.transaction_id "211 "JOIN risk_scores ON a.id=rs.account_id "212 "JOIN devices ON t.device_id=d.id "213 "JOIN locations ON t.location_id=loc.id "214 "WHERE t.amount > 10000 AND fl.is_flagged=true"215 ),216 tables=[217 TableSpec("locations", true_rows=1_000_000, selectivity=0.10, has_index=1, noise_sigma=0.10), # 100K218 TableSpec("fraud_labels", true_rows=3_000_000, selectivity=0.10, has_index=0, noise_sigma=0.40), # 300K219 TableSpec("merchants", true_rows=5_000_000, selectivity=0.10, has_index=1, noise_sigma=0.08), # 500K220 TableSpec("devices", true_rows=16_000_000, selectivity=0.05, has_index=1, noise_sigma=0.25), # 800K221 TableSpec("risk_scores", true_rows=12_000_000, selectivity=0.10, has_index=0, noise_sigma=0.30), # 1.2M222 TableSpec("accounts", true_rows=20_000_000, selectivity=0.08, has_index=1, noise_sigma=0.12), # 1.6M223 TableSpec("transactions", true_rows=50_000_000, selectivity=0.04, has_index=1, noise_sigma=0.25), # 2.0M224 ],225 ),226 Scenario(227 name="user_journey_attribution",228 domain="analytics",229 query_context=(230 "SELECT e.event_type, s.duration, u.segment, c.name, cv.revenue, ab.variant, pv.url "231 "FROM events JOIN sessions ON e.session_id=s.id "232 "JOIN users ON s.user_id=u.id "233 "JOIN campaigns ON s.campaign_id=c.id "234 "JOIN conversions ON s.id=cv.session_id "235 "JOIN ab_tests ON u.id=ab.user_id "236 "JOIN page_views ON s.id=pv.session_id "237 "WHERE s.started_at >= CURRENT_DATE - 7 AND c.status='active'"238 ),239 tables=[240 # Rebalanced: each table contributes cleanly to the cost range (no single dominator >30%)241 TableSpec("campaigns", true_rows=1_000_000, selectivity=0.10, has_index=1, noise_sigma=0.08), # 100K242 TableSpec("ab_tests", true_rows=3_000_000, selectivity=0.10, has_index=0, noise_sigma=0.30), # 300K243 TableSpec("users", true_rows=5_000_000, selectivity=0.10, has_index=1, noise_sigma=0.12), # 500K244 TableSpec("conversions", true_rows=16_000_000, selectivity=0.05, has_index=1, noise_sigma=0.20), # 800K245 TableSpec("sessions", true_rows=15_000_000, selectivity=0.08, has_index=1, noise_sigma=0.25), # 1.2M246 TableSpec("events", true_rows=32_000_000, selectivity=0.05, has_index=0, noise_sigma=0.50), # 1.6M247 TableSpec("page_views", true_rows=40_000_000, selectivity=0.05, has_index=0, noise_sigma=0.55), # 2.0M248 ],249 ),250]251 252# ─────────────────────────────────────────────────────────────────────────────253# Task Registry254# ─────────────────────────────────────────────────────────────────────────────255 256TASKS: dict[str, TaskConfig] = {257 "easy": TaskConfig(258 task_id="easy",259 name="OLTP Join Optimizer — 3-Table Queries",260 difficulty="easy",261 description=(262 "Production OLTP queries joining 3 tables from e-commerce and SaaS schemas. "263 "All indexes are available. Statistics are accurate. "264 "Goal: select the optimal join order and strategy. "265 "Scoring penalises nested-loop joins and missed indexes."266 ),267 scenarios=EASY_SCENARIOS,268 ),269 "medium": TaskConfig(270 task_id="medium",271 name="OLAP Join Optimizer — 5-Table Queries with Estimation Noise",272 difficulty="medium",273 description=(274 "Analytical queries joining 5 tables across e-commerce, marketing, "275 "and financial schemas. Some tables lack indexes. "276 "Cardinality estimates contain realistic noise (σ 0.05–0.25), "277 "requiring the agent to reason under uncertainty about true table sizes."278 ),279 scenarios=MEDIUM_SCENARIOS,280 ),281 "hard": TaskConfig(282 task_id="hard",283 name="Complex OLAP Join Optimizer — 7-Table Queries with High Estimation Noise",284 difficulty="hard",285 description=(286 "Enterprise analytical queries joining 7 tables with billion-row event "287 "tables, missing indexes, and high cardinality estimation noise (σ up to 0.60). "288 "Models real-world conditions: data skew, stale statistics, and column correlations "289 "that cause traditional planners to underestimate result sizes by 10–100×."290 ),291 scenarios=HARD_SCENARIOS,292 ),293}294 295 296# ─────────────────────────────────────────────────────────────────────────────297# Cardinality Estimation Noise298# ─────────────────────────────────────────────────────────────────────────────299 300def apply_estimation_noise(true_rows: int, noise_sigma: float, rng: random.Random) -> int:301 """302 Simulate cardinality estimation error using log-normal noise.303 304 real planners: estimated_rows = true_rows × exp(N(0, σ²))305 306 σ=0.05 → ~5% error (excellent statistics)307 σ=0.15 → ~15% error (typical OLTP)308 σ=0.25 → ~28% error (typical OLAP)309 σ=0.40 → ~49% error (stale statistics / large tables)310 σ=0.60 → ~82% error (data skew, column correlations)311 """312 if noise_sigma <= 0:313 return true_rows314 log_factor = rng.gauss(0.0, noise_sigma)315 estimated = int(true_rows * math.exp(log_factor))316 return max(1, estimated)317 318 319# ─────────────────────────────────────────────────────────────────────────────320# Cost Bound Computation (based on true_rows, not estimates)321# ─────────────────────────────────────────────────────────────────────────────322 323def compute_cost_bounds(tables: List[TableSpec]) -> tuple[float, float]:324 """325 Analytical worst/best cost bounds for a scenario, accounting for cascading sizes.326 Joining early inflates all subsequent step costs via an additive intermediate size penalty.327 328 Worst: worst ordering by size (largest first), nested-loop (2.0×), no index329 Best: best ordering by size (smallest first), merge-sort (0.8×), use index330 """331 # Order by intermediate output size to find best and worst sequences332 outputs = [(t.true_rows * t.selectivity, t) for t in tables]333 best_order = sorted(outputs, key=lambda x: x[0])334 worst_order = sorted(outputs, key=lambda x: x[0], reverse=True)335 336 def simulate_cost(order_list, is_worst: bool) -> float:337 total_cost = 0.0338 running_size = 0.0 # Cost penalty is 0 on the first table339 340 for output_size, t in order_list:341 if is_worst:342 step_cost = (t.true_rows * t.selectivity * 2.0) + running_size343 else:344 best_base = t.true_rows * 0.5 if t.has_index else t.true_rows345 step_cost = (best_base * t.selectivity * 0.8) + running_size346 347 total_cost += step_cost348 349 # For the first table, running_size becomes output_size. Afterwards it multiplies.350 if running_size == 0.0:351 running_size = output_size352 else:353 running_size *= output_size354 355 return total_cost356 357 worst = simulate_cost(worst_order, is_worst=True)358 best = simulate_cost(best_order, is_worst=False)359 return worst, best360 361 362def compute_step_cost_bounds(table: TableSpec, running_size: float = 1.0) -> tuple[float, float]:363 """Worst/best cost bounds for a single table step, given additive penalty."""364 # Note: running_size defaults to 1.0 in env (or actual size later), we subtract 1.0 to behave like 0.0 on step 1365 penalty = max(0.0, running_size - 1.0)366 worst = (table.true_rows * table.selectivity * 2.0) + penalty367 best_base = table.true_rows * 0.5 if table.has_index else table.true_rows368 best = (best_base * table.selectivity * 0.8) + penalty369 return worst, best370 371 372# ─────────────────────────────────────────────────────────────────────────────373# Order Quality Scorer — measures how well the join sequence minimises374# intermediate result explosion (independent of join algorithm choices)375# ─────────────────────────────────────────────────────────────────────────────376 377def compute_order_quality(chosen_order: List[int], tables: List[TableSpec]) -> float:378 """379 Score the join ordering quality using log-normalised rank mismatch.380 381 Optimal order: join tables with smallest (true_rows × selectivity) first,382 because this minimises intermediate result sizes throughout the sequence.383 Joining a large-output table early compounds the cost of every subsequent384 join — this is the core of real query optimisation difficulty.385 386 Method:387 1. Compute each table's output size = true_rows × selectivity388 2. For each position k, compare the log-output of the chosen k-th table389 against the log-output of the k-th optimal table (sorted ascending)390 3. Normalise mismatch against the worst possible ordering (descending)391 392 Returns 1.0 for a perfectly sorted order, 0.0 for reverse optimal.393 """394 n = len(tables)395 if n <= 1:396 return 1.0397 398 outputs = [t.true_rows * t.selectivity for t in tables]399 chosen_outputs = [outputs[i] for i in chosen_order]400 optimal_outputs = sorted(outputs) # ascending — best ordering401 worst_outputs = list(reversed(optimal_outputs)) # descending — worst ordering402 403 def log_mismatch(order_outputs: List[float]) -> float:404 return sum(405 abs(math.log1p(order_outputs[k]) - math.log1p(optimal_outputs[k]))406 for k in range(n)407 )408 409 actual_mismatch = log_mismatch(chosen_outputs)410 max_mismatch = log_mismatch(worst_outputs)411 412 if max_mismatch < 1e-9:413 return 1.0414 415 score = 1.0 - actual_mismatch / max_mismatch416 return float(max(0.0, min(1.0, score)))417 418 419# ─────────────────────────────────────────────────────────────────────────────420# Grader — returns float in [0.0, 1.0]421# ─────────────────────────────────────────────────────────────────────────────422 423def grade(424 tables: List[TableSpec],425 final_cost: float,426 chosen_order: Optional[List[int]] = None,427) -> float:428 """429 Grade a completed episode. Combined score = 60% method + 40% order.430 431 method_score: how well the agent chose join algorithms and index usage.432 = (worst_cost - actual_cost) / (worst_cost - best_cost)433 worst = all nested-loop, no indexes434 best = all merge-sort, use index when available435 436 order_score: how well the agent sequenced the joins.437 Optimal = smallest-output tables first (minimises intermediate blowup).438 Computed via log-normalised rank mismatch vs optimal ordering.439 Only applied when all tables have been joined (full episode).440 441 Combined = 0.60 × method_score + 0.40 × order_score442 443 This ensures:444 - A model that always picks merge_sort but joins biggest tables first445 scores ~0.60 instead of ~1.00 on hard tasks.446 - A model with correct ordering but mediocre method choices still447 receives partial credit from the order component.448 - Difficulty genuinely increases: easy (3 tables, 6 orderings) vs449 hard (7 tables, 5040 orderings) differentiates planning depth.450 """451 # ── Method quality (join algorithm + index usage) ─────────────────────────452 worst, best = compute_cost_bounds(tables)453 if worst == best:454 method_score = 1.0455 else:456 method_score = float(max(0.0, min(1.0, (worst - final_cost) / (worst - best))))457 458 # ── Order quality (join sequence planning) ────────────────────────────────459 if chosen_order is not None and len(chosen_order) == len(tables):460 order_score = compute_order_quality(chosen_order, tables)461 else:462 # Incomplete episode or no order info — fall back to method score only463 order_score = method_score464 465 # ── Combined score ────────────────────────────────────────────────────────466 combined = 0.60 * method_score + 0.40 * order_score467 return float(max(0.0, min(1.0, combined)))468