RonyForAI/Mirage_DB_RL

1---2title: Mirage RL — Production Query Optimizer Environment3emoji: 🗄️4colorFrom: blue5colorTo: indigo6sdk: docker7pinned: false8app_port: 80009tags:10  - openenv11  - reinforcement-learning12  - database13  - query-optimization14---15 16# Mirage_RL: Production Query Join-Order Optimizer17 18> **A reinforcement learning environment for the #1 performance bottleneck in production databases — join-order optimisation under cardinality estimation uncertainty.**19 20---21 22## Why This Problem Is Real and Hard23 24Every analytical database — PostgreSQL, MySQL, Snowflake, BigQuery, Spark SQL — must decide **how to join tables** before executing a query. For a query touching 7 tables, there are **5,040 possible join orders**. The difference between the best and worst order can be **100–1,000× in execution time**.25 26The challenge is that real planners don't know the true intermediate result sizes. They rely on table statistics (histograms, row counts, distinct-value estimates) that are frequently:27- **Stale** — vacuumed days or weeks ago on large tables28- **Wrong under data skew** — top 1% of customers produce 80% of orders29- **Compounding** — errors multiply across joins (the "selectivity estimation problem")30 31This causes PostgreSQL and MySQL to routinely choose plans that are **10–100× slower** than optimal for complex analytical queries. It is a **known, open, expensive problem** — teams at [Meta](https://research.facebook.com/), [Snowflake](https://www.snowflake.com/), [Databricks](https://databricks.com/), and [Carnegie Mellon](https://vldb.org/pvldb/vol12/p1705-marcus.pdf) actively work on learned query optimizers.32 33**Mirage_RL** provides a rigorous RL benchmark for this exact problem. An AI agent must learn to plan optimal join orders **under the same uncertainty** a real database planner faces.34 35---36 37## Environment Overview38 39Mirage_RL simulates the **planning phase** of a query optimizer. At each step, the agent decides:40 411. **Which table** to join next (from the remaining unjoined tables)422. **Which join algorithm** to use — hash (×1.0), nested-loop (×2.0, avoid), or merge-sort (×0.8)433. **Whether to use an index** (halves base row count when a covering index exists)44 45The agent sees **noisy cardinality estimates** — not ground truth — exactly matching real planner conditions. The actual execution cost is computed on true row counts (hidden from the agent), creating a realistic planning-under-uncertainty problem.46 47**Episode domains** are sampled from three production enterprise schemas:48| Domain | Tables | Typical scale |49|---|---|---|50| **E-commerce (OLTP)** | orders, customers, products, categories, suppliers, warehouses | 10K – 50M rows |51| **Analytics (OLAP)** | events, sessions, users, campaigns, conversions, ab_tests | 1M – 1B rows |52| **Financial (OLAP)** | transactions, accounts, merchants, fraud_labels, risk_scores | 5M – 100M rows |53 54---55 56## Action Space57 58```python59class QueryAction(Action):60    next_table: int   # Index of next table to join (must be from remaining_tables)61    join_type:  int   # 0=hash(×1.0)  1=nested-loop(×2.0)  2=merge-sort(×0.8)62    use_index:  int   # 0=full scan   1=index scan (halves base rows if index exists)63```64 65**Cost model:**66```67step_cost  =  base_rows  ×  selectivity  ×  join_multiplier68base_rows  =  estimated_rows × 0.5   if use_index=1 AND index present69           =  estimated_rows          otherwise70```71 72An optimal agent should **always prefer merge-sort** and **always use an index when available**.73 74---75 76## Observation Space77 78```python79class QueryObservation(Observation):80    tables:           List[str]   # table names in this query81    table_rows:       List[int]   # estimated row counts (noisy — simulates planner statistics)82    selectivities:    List[float] # join predicate selectivity per table (0.0–1.0)83    has_index:        List[int]   # index availability: 1=yes, 0=no84    chosen_order:     List[int]   # indices of tables already joined85    remaining_tables: List[int]   # indices of tables not yet joined86    step_number:      int         # 0-based step counter for this episode87    current_cost:     float       # accumulated join cost so far (true rows)88    query_context:    str         # SQL-style description of the query being optimized89```90 91> **⚠️ Cardinality estimation noise.** `table_rows` in the observation reflects **estimated** cardinalities, not ground truth — matching real-world conditions where planners operate on stale statistics. Noise sigma varies by task difficulty:92> - Easy: σ ≈ 0.03–0.06 (excellent statistics)93> - Medium: σ ≈ 0.05–0.25 (typical production OLTP/OLAP)94> - Hard: σ ≈ 0.10–0.60 (stale stats, billion-row event tables, data skew)95 96---97 98## Task Definitions99 100### Task 1: OLTP Join Optimizer — Easy (3 tables)101 102**Scenarios:** E-commerce catalog lookup, SaaS subscription query, inventory reorder check103 104All tables have covering indexes. Statistics are fresh and accurate. The agent must learn:105- Join smaller dimension tables (categories, suppliers) before larger fact tables106- Always prefer merge-sort joins107- Use indexes to reduce base row counts108 109**Scoring:** `score = (worst_cost - actual_cost) / (worst_cost - best_cost)` in [0.0, 1.0]110 111Expected baseline (random agent): ~0.35 | Expected upper bound (optimal): ~0.95+112 113---114 115### Task 2: OLAP Join Optimizer — Medium (5 tables)116 117**Scenarios:** E-commerce order fulfillment, marketing funnel analytics, financial transaction summary118 119Mix of indexed and unindexed tables. Cardinality estimates have realistic noise (σ up to 0.25). The agent must simultaneously:120- Identify efficient join orderings by reasoning about **estimated** table sizes121- Handle missing indexes by choosing merge-sort over nested-loop122- Navigate 5! = 120 possible orderings123 124Expected baseline (random agent): ~0.25 | Expected upper bound (optimal): ~0.85+125 126---127 128### Task 3: Complex OLAP Join Optimizer — Hard (7 tables)129 130**Scenarios:** Full e-commerce pipeline audit, financial fraud detection, user journey attribution131 132Seven-table joins with billion-row event tables, missing indexes, and high estimation noise (σ up to 0.60). The optimal join order requires understanding which tables filter aggressively via selectivity, which have indexes, and which are so large they must come last.133 134- 7! = **5,040 possible join orderings**135- Estimation errors up to ±80% on raw row counts136- Three join algorithms × two index options per step = 6 action variants per table137 138Expected baseline (random agent): ~0.15 | Expected upper bound (frontier LLM): ~0.70+139 140---141 142## Reward Function143 144Rewards are **normalised to [0.0, 1.0]** at every step — never sparse, never binary.145 146```147Per-step reward = (worst_step_cost - actual_step_cost) / (worst_step_cost - best_step_cost)148 149Final reward    = (worst_total_cost - actual_total_cost) / (worst_total_cost - best_total_cost)150```151 152| Scenario | Reward |153|---|---|154| Optimal join choice (merge-sort + index) | 1.0 |155| Hash join, index used | ~0.7–0.9 |156| Merge-sort, no index (table has none) | 1.0 |157| Hash join, ignores available index | 0.4–0.6 |158| Nested-loop join | 0.0 |159| Invalid table selection (already joined) | 0.0 |160 161This design gives the agent a **learning signal at every step**, enabling effective credit assignment for RL from any trajectory.162 163---164 165## Setup & Usage166 167### Prerequisites168 169```bash170pip install openenv-core openai171```172 173### Run the Server Locally174 175```bash176# From the Mirage_RL/ directory177uvicorn server.app:app --host 0.0.0.0 --port 8000178```179 180### Run with Docker181 182```bash183# Build184docker build -t mirage-rl:latest .185 186# Run187docker run -p 8000:8000 mirage-rl:latest188```189 190### API Endpoints191 192| Endpoint | Method | Description |193|---|---|---|194| `/reset` | POST | Start new episode. Body: `{"task_id": "easy\|medium\|hard", "seed": 42}` |195| `/step` | POST | Execute action. Body: `{"next_table": 0, "join_type": 2, "use_index": 1}` |196| `/state` | GET | Get current environment state |197| `/health` | GET | Health check |198| `/docs` | GET | OpenAPI documentation |199 200### Run Baseline Inference Script201 202```bash203# Set required environment variables204export HF_TOKEN=your_huggingface_token205export API_BASE_URL=https://router.huggingface.co/v1206export MODEL_NAME=Qwen/Qwen2.5-72B-Instruct207 208# Run against all 3 tasks209python inference.py210```211 212### Python Client213 214```python215from Mirage_RL import QueryClient, QueryAction216 217with QueryClient(base_url="http://localhost:8000").sync() as env:218    result = env.reset()                 # starts easy task by default219    obs = result.observation220    print(f"Query: {obs.query_context}")221    print(f"Tables: {obs.tables}")222 223    result = env.step(QueryAction(224        next_table=0,   # join first table225        join_type=2,    # merge-sort226        use_index=1,    # use index227    ))228    print(f"Reward: {result.reward:.4f}")229```230 231---232 233## Baseline Scores234 235Scores produced by running `inference.py` with `Qwen/Qwen2.5-72B-Instruct` via HuggingFace Inference Router:236 237| Task | Difficulty | Baseline Score |238|---|---|---|239| OLTP 3-table | Easy | ~1.00 |240| OLAP 5-table | Medium | ~0.93 |241| Complex OLAP 7-table | Hard | ~0.11 |242 243> Scores are approximate and vary by episode (random scenario sampling). For reproducible results, pass `seed=42` to `reset()`.244 245---246 247## Project Structure248 249```250Mirage_RL/251├── Dockerfile                         # Container image (root-level for HF Spaces)252├── README.md                          # This file253├── openenv.yaml                       # OpenEnv manifest (spec_version, tasks, schemas)254├── pyproject.toml                     # Project metadata and dependencies255├── inference.py                       # Baseline inference script (mandatory)256├── models.py                          # Pydantic models: QueryAction, QueryObservation, QueryState257├── client.py                          # QueryClient: HTTP client for the server258├── __init__.py                        # Package exports259└── server/260    ├── app.py                         # FastAPI application (HTTP endpoints)261    ├── Mirage_RL_environment.py       # Core environment logic (reset/step/state)262    ├── tasks.py                       # Enterprise scenario definitions + graders263    ├── Dockerfile                     # Docker build (also accessible from server/)264    └── requirements.txt               # Minimal server dependencies265```266 267---268 269## Validation270 271Run the OpenEnv pre-submission validator:272 273```bash274# From the Mirage_RL/ directory275openenv validate276 277# Full submission validator (requires deployed HF Space URL)278bash validate-submission.sh https://your-space.hf.space .279```280 281---282 283## References284 285- [Neo: A Learned Query Optimizer](https://vldb.org/pvldb/vol12/p1705-marcus.pdf) — Marcus et al., VLDB 2019286- [Bao: Making Learned Query Optimization Practical](https://dl.acm.org/doi/10.1145/3448016.3452838) — Marcus et al., SIGMOD 2021287- [Are We Ready for Learned Cardinality Estimation?](https://vldb.org/pvldb/vol14/p1640-wang.pdf) — Wang et al., VLDB 2021288- [PostgreSQL Query Planner](https://www.postgresql.org/docs/current/planner-optimizer.html) — Production optimizer documentation289