From b1776b17f4634fdc6d38a7610e9e5e5b69cb2903 Mon Sep 17 00:00:00 2001 From: OfficialAbhinavSingh Date: Thu, 23 Jul 2026 20:18:22 +0530 Subject: [PATCH 1/4] feat(sql-optim-env): add execution-grounded DuckDB SQL query optimization environment --- docs/source/_toctree.yml | 2 + docs/source/environments.md | 7 + docs/source/environments/sql_optim.md | 95 ++++ envs/sql_optim_env/README.md | 110 ++++ envs/sql_optim_env/__init__.py | 34 ++ envs/sql_optim_env/client.py | 72 +++ envs/sql_optim_env/models.py | 123 ++++ envs/sql_optim_env/openenv.yaml | 6 + envs/sql_optim_env/pyproject.toml | 37 ++ envs/sql_optim_env/server/Dockerfile | 55 ++ envs/sql_optim_env/server/__init__.py | 11 + envs/sql_optim_env/server/app.py | 39 ++ envs/sql_optim_env/server/executor.py | 324 +++++++++++ envs/sql_optim_env/server/graders.py | 219 ++++++++ envs/sql_optim_env/server/scoring_models.py | 43 ++ .../server/sql_optim_environment.py | 194 +++++++ envs/sql_optim_env/server/tasks.py | 523 ++++++++++++++++++ tests/envs/test_sql_optim_environment.py | 384 +++++++++++++ 18 files changed, 2278 insertions(+) create mode 100644 docs/source/environments/sql_optim.md create mode 100644 envs/sql_optim_env/README.md create mode 100644 envs/sql_optim_env/__init__.py create mode 100644 envs/sql_optim_env/client.py create mode 100644 envs/sql_optim_env/models.py create mode 100644 envs/sql_optim_env/openenv.yaml create mode 100644 envs/sql_optim_env/pyproject.toml create mode 100644 envs/sql_optim_env/server/Dockerfile create mode 100644 envs/sql_optim_env/server/__init__.py create mode 100644 envs/sql_optim_env/server/app.py create mode 100644 envs/sql_optim_env/server/executor.py create mode 100644 envs/sql_optim_env/server/graders.py create mode 100644 envs/sql_optim_env/server/scoring_models.py create mode 100644 envs/sql_optim_env/server/sql_optim_environment.py create mode 100644 envs/sql_optim_env/server/tasks.py create mode 100644 tests/envs/test_sql_optim_environment.py diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 499138c66..20235d45a 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -131,6 +131,8 @@ title: OpenCode - local: environments/sophistry_bench_sprint title: Sophistry Bench Sprint + - local: environments/sql_optim + title: SQL Query Optimization title: Environments - isExpanded: false sections: diff --git a/docs/source/environments.md b/docs/source/environments.md index beb19be62..2dc7cfd89 100644 --- a/docs/source/environments.md +++ b/docs/source/environments.md @@ -265,6 +265,13 @@ The OpenEnv community has built a catalog of ready-to-run environments that cove šŸ“„ Docs +
+
SQL Query Optimization
+

sql_optim_env rewards an agent for rewriting slow SQL: it executes both the original and optimized query against an in-memory DuckDB seeded with synthetic data and scores measured speedup plus result-correctness across five anti-pattern tasks.

+
+ šŸ“„ Docs +
+
diff --git a/docs/source/environments/sql_optim.md b/docs/source/environments/sql_optim.md new file mode 100644 index 000000000..7e2529a7b --- /dev/null +++ b/docs/source/environments/sql_optim.md @@ -0,0 +1,95 @@ + +# SQL Query Optimization Environment + +An **execution-grounded** environment for training agents to write fast SQL. The +agent is given a slow query plus its schema and must return a rewrite. Reward is +**not** a keyword heuristic: the environment runs both the original and the +optimized query against a real in-memory **DuckDB** database seeded with +realistic synthetic data (10k users, 500k orders, 1k products, 1M events), then +scores the agent on measured speedup and result-correctness. + +Five tasks scale from basic anti-patterns to expert window-function audits. + +## Quick Start + +```python +from sql_optim_env import SQLOptimAction, SQLOptimEnv + +with SQLOptimEnv(base_url="http://localhost:8000").sync() as env: + obs = env.reset(task_id="task_1_basic_antipatterns").observation + print(obs.sql_query) # the slow query to fix + print(obs.schema_info) # tables, row counts + + result = env.step( + SQLOptimAction( + suggestions=[ + {"issue_type": "select_star", "line": 1, + "description": "SELECT * reads every column", "severity": "high", + "fix": "project only the needed columns"}, + ], + optimized_query="SELECT id FROM orders WHERE customer_id = 5000", + summary="Dropped SELECT * and the non-sargable CAST.", + estimated_improvement="~10x faster", + approved=False, + ) + ) + print(result.reward, result.done) + print(result.metadata["reward_breakdown"]) + print(result.metadata["execution"]["speedup"]) # real DuckDB speedup +``` + +Async usage mirrors the other environments: + +```python +async with SQLOptimEnv(base_url="http://localhost:8000") as env: + result = await env.reset(task_id="task_1_basic_antipatterns") +``` + +## Action + +| Field | Type | Meaning | +|-------|------|---------| +| `suggestions` | `list[dict]` | Issues found: `issue_type`, `line`, `description`, `severity`, `fix` | +| `optimized_query` | `str` | The rewritten SQL ā€” executed for real | +| `summary` | `str` | Overall analysis | +| `estimated_improvement` | `str` | Agent's own speedup estimate | +| `approved` | `bool` | `True` if the query is already optimal | + +## Observation + +`sql_query`, `schema_info`, `task_description`, `difficulty`, `step_count`, +`max_steps`, `issues_found_so_far`, and `last_execution` (the previous step's +real timing comparison, so the agent can refine its rewrite). The step reward is +on `reward`; the composite `reward_breakdown`, `feedback`, and full `execution` +comparison are under `metadata`. + +## Reward + +Composite score in `[0, 1]`, computed inside the environment: + +| Component | Weight | Source | +|-----------|--------|--------| +| Execution speedup | 35% | real DuckDB timing ratio | +| Result correctness | 20% | do both queries return identical rows? | +| Issue detection | 25% | flagged issues vs. ground truth | +| Approval correctness | 8% | correctly judged the query bad/good | +| Summary quality | 7% | thoroughness of the written analysis | +| Severity labels | 5% | severities present | + +## Tasks + +Five tasks (`task_1_basic_antipatterns` ā€¦ `task_5_*`) covering `SELECT *`, +non-sargable predicates, functions on filter columns, join ordering, and window +functions, spanning `easy` ā†’ `expert`. + +## Run the server locally + +```bash +uvicorn server.app:app --host 0.0.0.0 --port 8000 +``` + +## Tests + +```bash +PYTHONPATH=src:envs uv run pytest tests/envs/test_sql_optim_environment.py -v +``` diff --git a/envs/sql_optim_env/README.md b/envs/sql_optim_env/README.md new file mode 100644 index 000000000..20d72e633 --- /dev/null +++ b/envs/sql_optim_env/README.md @@ -0,0 +1,110 @@ +--- +title: SQL Query Optimization Environment +emoji: šŸ—„ļø +colorFrom: indigo +colorTo: blue +sdk: docker +pinned: false +app_port: 8000 +base_path: /web +tags: + - openenv + - sql + - duckdb + - optimization +--- + +# SQL Query Optimization Environment + +An **execution-grounded** environment for training agents to write fast SQL. The +agent is given a slow query plus its schema and must return a rewrite. Reward is +**not** a keyword heuristic: the environment runs both the original and the +optimized query against a real in-memory **DuckDB** database seeded with +realistic synthetic data (10k users, 500k orders, 1k products, 1M events), then +scores the agent on measured speedup and result-correctness. + +Five tasks scale from basic anti-patterns to expert window-function audits. + +## Quick Start + +```python +from sql_optim_env import SQLOptimAction, SQLOptimEnv + +with SQLOptimEnv(base_url="http://localhost:8000").sync() as env: + obs = env.reset(task_id="task_1_basic_antipatterns").observation + print(obs.sql_query) # the slow query to fix + print(obs.schema_info) # tables, row counts + + result = env.step( + SQLOptimAction( + suggestions=[ + {"issue_type": "select_star", "line": 1, + "description": "SELECT * reads every column", "severity": "high", + "fix": "project only the needed columns"}, + ], + optimized_query="SELECT id FROM orders WHERE customer_id = 5000", + summary="Dropped SELECT * and the non-sargable CAST.", + estimated_improvement="~10x faster", + approved=False, + ) + ) + print(result.reward, result.done) + print(result.metadata["reward_breakdown"]) + print(result.metadata["execution"]["speedup"]) # real DuckDB speedup +``` + +Async usage mirrors the other environments: + +```python +async with SQLOptimEnv(base_url="http://localhost:8000") as env: + result = await env.reset(task_id="task_1_basic_antipatterns") +``` + +## Action + +| Field | Type | Meaning | +|-------|------|---------| +| `suggestions` | `list[dict]` | Issues found: `issue_type`, `line`, `description`, `severity`, `fix` | +| `optimized_query` | `str` | The rewritten SQL ā€” executed for real | +| `summary` | `str` | Overall analysis | +| `estimated_improvement` | `str` | Agent's own speedup estimate | +| `approved` | `bool` | `True` if the query is already optimal | + +## Observation + +`sql_query`, `schema_info`, `task_description`, `difficulty`, `step_count`, +`max_steps`, `issues_found_so_far`, and `last_execution` (the previous step's +real timing comparison, so the agent can refine its rewrite). The step reward is +on `reward`; the composite `reward_breakdown`, `feedback`, and full `execution` +comparison are under `metadata`. + +## Reward + +Composite score in `[0, 1]`, computed inside the environment: + +| Component | Weight | Source | +|-----------|--------|--------| +| Execution speedup | 35% | real DuckDB timing ratio | +| Result correctness | 20% | do both queries return identical rows? | +| Issue detection | 25% | flagged issues vs. ground truth | +| Approval correctness | 8% | correctly judged the query bad/good | +| Summary quality | 7% | thoroughness of the written analysis | +| Severity labels | 5% | severities present | + +## Tasks + +Five tasks (`task_1_basic_antipatterns` ā€¦ `task_5_*`) covering `SELECT *`, +non-sargable predicates, functions on filter columns, join ordering, and window +functions, spanning `easy` ā†’ `expert`. + +## Run the server locally + +```bash +uvicorn server.app:app --host 0.0.0.0 --port 8000 +``` + +## Tests + +```bash +PYTHONPATH=src:envs uv run pytest tests/envs/test_sql_optim_environment.py -v +``` diff --git a/envs/sql_optim_env/__init__.py b/envs/sql_optim_env/__init__.py new file mode 100644 index 000000000..71a4f4450 --- /dev/null +++ b/envs/sql_optim_env/__init__.py @@ -0,0 +1,34 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +SQL Query Optimization Environment for OpenEnv. + +An execution-grounded RL environment: the agent rewrites slow SQL and is +rewarded by real DuckDB execution speedup plus result-correctness across five +tasks that scale from basic anti-patterns to expert window-function audits. + +Examples: + +```python +from envs.sql_optim_env import SQLOptimEnv, SQLOptimAction + +with SQLOptimEnv(base_url="http://localhost:8000").sync() as env: + obs = env.reset(task_id="task_1_basic_antipatterns").observation + result = env.step(SQLOptimAction(optimized_query="SELECT id FROM orders ...")) + print(result.reward, result.done) +``` +""" + +from .client import SQLOptimEnv +from .models import SQLOptimAction, SQLOptimObservation, SQLOptimState + +__all__ = [ + "SQLOptimEnv", + "SQLOptimAction", + "SQLOptimObservation", + "SQLOptimState", +] diff --git a/envs/sql_optim_env/client.py b/envs/sql_optim_env/client.py new file mode 100644 index 000000000..0ac468319 --- /dev/null +++ b/envs/sql_optim_env/client.py @@ -0,0 +1,72 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Client for the SQL Query Optimization environment. + +Maintains a persistent WebSocket connection to the environment server for +efficient multi-step episodes. Async by default; use `.sync()` for a +synchronous wrapper. +""" + +from __future__ import annotations + +from typing import Any, Dict + +from openenv.core.client_types import StepResult +from openenv.core.env_client import EnvClient + +from .models import SQLOptimAction, SQLOptimObservation, SQLOptimState + + +class SQLOptimEnv(EnvClient[SQLOptimAction, SQLOptimObservation, SQLOptimState]): + """ + Client for the SQL Query Optimization environment. + + The agent receives a slow SQL query plus its schema, returns a + [`SQLOptimAction`][envs.sql_optim_env.models.SQLOptimAction] with a + rewritten `optimized_query`, and is rewarded by real DuckDB execution + speedup and result-correctness. + + Examples: + + ```python + with SQLOptimEnv(base_url="http://localhost:8000").sync() as env: + obs = env.reset(task_id="task_1_basic_antipatterns").observation + result = env.step( + SQLOptimAction( + optimized_query="SELECT id FROM orders WHERE customer_id = 5000", + summary="Dropped SELECT * and the non-sargable CAST.", + ) + ) + print(result.reward, result.done) + ``` + """ + + def _step_payload(self, action: SQLOptimAction) -> Dict[str, Any]: + """Serialize a [`SQLOptimAction`] into the step request payload.""" + return { + "suggestions": action.suggestions, + "optimized_query": action.optimized_query, + "summary": action.summary, + "estimated_improvement": action.estimated_improvement, + "approved": action.approved, + } + + def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SQLOptimObservation]: + """Parse a server step/reset response into a `StepResult`.""" + obs_data = payload.get("observation", {}) + observation = SQLOptimObservation(**obs_data) + return StepResult( + observation=observation, + reward=payload.get("reward", observation.reward), + done=payload.get("done", observation.done), + metadata=payload.get("metadata", observation.metadata), + ) + + def _parse_state(self, payload: Dict[str, Any]) -> SQLOptimState: + """Parse a server state response into a [`SQLOptimState`].""" + return SQLOptimState(**payload) diff --git a/envs/sql_optim_env/models.py b/envs/sql_optim_env/models.py new file mode 100644 index 000000000..d8cfdd413 --- /dev/null +++ b/envs/sql_optim_env/models.py @@ -0,0 +1,123 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Data models for the SQL Query Optimization environment. + +Defines the `Action`, `Observation`, and `State` wire types exchanged between +the [`SQLOptimEnv`][envs.sql_optim_env.client.SQLOptimEnv] client and the +environment server. All three inherit from the OpenEnv base types so they +serialize over the WebSocket/HTTP transport. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from openenv.core.env_server import Action, Observation, State +from pydantic import Field + + +class SQLOptimAction(Action): + """ + Action emitted by the agent for one optimization step. + + Attributes: + suggestions (`list[dict]`): + Detected issues. Each entry is a dict with `issue_type`, `line`, + `description`, `severity`, and `fix` keys. + optimized_query (`str`): + The complete rewritten SQL. It is executed against the real DuckDB + data so the measured speedup can be rewarded. + summary (`str`): + Overall analysis and performance profile of the query. + estimated_improvement (`str`): + The agent's own expected speedup (e.g. `"10x faster"`). + approved (`bool`): + `True` if the agent judges the query already optimal, `False` if it + still needs changes. + """ + + suggestions: List[Dict[str, Any]] = Field(default_factory=list) + optimized_query: str = "" + summary: str = "" + estimated_improvement: str = "" + approved: bool = False + + +class SQLOptimObservation(Observation): + """ + Observation returned to the agent after `reset` and each `step`. + + Inherits `done`, `reward`, and `metadata` from + [`~openenv.core.env_server.Observation`]. The composite reward breakdown, + per-criterion feedback, and full execution comparison are attached under + `metadata` (keys `reward_breakdown`, `feedback`, `execution`). + + Attributes: + task_id (`str`): + Identifier of the active task. + task_name (`str`): + Human-readable task name. + task_description (`str`): + What the agent must do for this task. + sql_query (`str`): + The slow SQL query to analyze and rewrite. + schema_info (`str`): + Table schema, row counts, and index notes. + dialect (`str`): + SQL dialect the query targets. + difficulty (`str`): + One of `easy`, `medium`, `hard`, `expert`. + step_count (`int`): + Steps taken so far in this episode. + max_steps (`int`): + Maximum steps allowed for this task. + issues_found_so_far (`list[str]`): + Issue types the agent flagged in previous steps. + last_execution (`dict`, *optional*): + Execution comparison from the previous step, so the agent can refine + its `optimized_query`. + """ + + task_id: str = "" + task_name: str = "" + task_description: str = "" + sql_query: str = "" + schema_info: str = "" + dialect: str = "duckdb/postgresql" + difficulty: str = "" + step_count: int = 0 + max_steps: int = 5 + issues_found_so_far: List[str] = Field(default_factory=list) + last_execution: Optional[Dict[str, Any]] = None + + +class SQLOptimState(State): + """ + Internal environment state for the current episode. + + Inherits `episode_id` and `step_count` from + [`~openenv.core.env_server.State`]. + + Attributes: + task_id (`str`): + Identifier of the active task (`"none"` when no episode is active). + max_steps (`int`): + Maximum steps allowed for the active task. + episode_done (`bool`): + Whether the current episode has terminated. + cumulative_reward (`float`): + Sum of per-step rewards in this episode. + current_task (`str`): + Human-readable name of the active task. + """ + + task_id: str = "none" + max_steps: int = 0 + episode_done: bool = True + cumulative_reward: float = 0.0 + current_task: str = "No active episode" diff --git a/envs/sql_optim_env/openenv.yaml b/envs/sql_optim_env/openenv.yaml new file mode 100644 index 000000000..2aff14057 --- /dev/null +++ b/envs/sql_optim_env/openenv.yaml @@ -0,0 +1,6 @@ +spec_version: 1 +name: sql_optim_env +type: space +runtime: fastapi +app: server.app:app +port: 8000 diff --git a/envs/sql_optim_env/pyproject.toml b/envs/sql_optim_env/pyproject.toml new file mode 100644 index 000000000..7820ff355 --- /dev/null +++ b/envs/sql_optim_env/pyproject.toml @@ -0,0 +1,37 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +[build-system] +requires = ["setuptools>=45", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "openenv-sql-optim-env" +version = "0.1.0" +description = "SQL Query Optimization Environment for OpenEnv - execution-grounded RL with real DuckDB timing" +requires-python = ">=3.10" +dependencies = [ + "openenv>=0.2.3", + "fastapi>=0.115.0", + "pydantic>=2.0.0", + "uvicorn>=0.24.0", + "requests>=2.31.0", + "duckdb>=0.10.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0.0", + "pytest-cov>=4.0.0", +] + +[project.scripts] +server = "sql_optim_env.server.app:main" + +[tool.setuptools] +include-package-data = true +packages = ["sql_optim_env", "sql_optim_env.server"] +package-dir = { "sql_optim_env" = ".", "sql_optim_env.server" = "server" } diff --git a/envs/sql_optim_env/server/Dockerfile b/envs/sql_optim_env/server/Dockerfile new file mode 100644 index 000000000..85aa71064 --- /dev/null +++ b/envs/sql_optim_env/server/Dockerfile @@ -0,0 +1,55 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest +FROM ${BASE_IMAGE} AS builder + +WORKDIR /app + + +COPY . /app/env + +WORKDIR /app/env + +RUN if ! command -v uv >/dev/null 2>&1; then \ + curl -LsSf https://astral.sh/uv/install.sh | sh && \ + mv /root/.local/bin/uv /usr/local/bin/uv && \ + mv /root/.local/bin/uvx /usr/local/bin/uvx; \ + fi + +RUN apt-get update && apt-get install -y --no-install-recommends \ + git \ + && rm -rf /var/lib/apt/lists/* + +RUN --mount=type=cache,target=/root/.cache/uv \ + if [ -f uv.lock ]; then \ + uv sync --frozen --no-install-project --no-editable; \ + else \ + uv sync --no-install-project --no-editable; \ + fi + +RUN --mount=type=cache,target=/root/.cache/uv \ + if [ -f uv.lock ]; then \ + uv sync --frozen --no-editable; \ + else \ + uv sync --no-editable; \ + fi + +# Final runtime stage +FROM ${BASE_IMAGE} + +WORKDIR /app + +COPY --from=builder /app/env/.venv /app/.venv +COPY --from=builder /app/env /app/env + +ENV PATH="/app/.venv/bin:$PATH" +ENV PYTHONPATH="/app/env:$PYTHONPATH" + +HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \ + CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1 + +CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"] diff --git a/envs/sql_optim_env/server/__init__.py b/envs/sql_optim_env/server/__init__.py new file mode 100644 index 000000000..00ea1c181 --- /dev/null +++ b/envs/sql_optim_env/server/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""SQL Query Optimization environment server.""" + +from .sql_optim_environment import SQLOptimEnvironment + +__all__ = ["SQLOptimEnvironment"] diff --git a/envs/sql_optim_env/server/app.py b/envs/sql_optim_env/server/app.py new file mode 100644 index 000000000..9472cac7a --- /dev/null +++ b/envs/sql_optim_env/server/app.py @@ -0,0 +1,39 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""FastAPI application for the SQL Query Optimization environment.""" + +from openenv.core.env_server import create_app + +# Support both in-repo and standalone imports +try: + # In-repo imports (running from the OpenEnv repository) + from ..models import SQLOptimAction, SQLOptimObservation + from .sql_optim_environment import SQLOptimEnvironment +except ImportError as e: + if "relative import" not in str(e) and "no known parent package" not in str(e): + raise + # Standalone imports (running via uvicorn server.app:app) + from models import SQLOptimAction, SQLOptimObservation + from server.sql_optim_environment import SQLOptimEnvironment + +# Pass the class (factory) rather than an instance for WebSocket session support. +app = create_app( + SQLOptimEnvironment, + SQLOptimAction, + SQLOptimObservation, + env_name="sql_optim_env", +) + + +def main(): + import uvicorn + + uvicorn.run(app, host="0.0.0.0", port=8000) + + +if __name__ == "__main__": + main() diff --git a/envs/sql_optim_env/server/executor.py b/envs/sql_optim_env/server/executor.py new file mode 100644 index 000000000..bc543dcc6 --- /dev/null +++ b/envs/sql_optim_env/server/executor.py @@ -0,0 +1,324 @@ +""" +executor.py ā€” DuckDB In-Memory SQL Execution Engine +===================================================== +The core innovation of this environment: instead of keyword-matching +heuristics, we ACTUALLY execute both the original and optimized queries +against realistic synthetic data and measure real performance differences. + +Tables populated: + users ā€” 10,000 rows + orders ā€” 500,000 rows + products ā€” 1,000 rows + events ā€” 1,000,000 rows +""" + +import atexit +import os +import shutil +import tempfile +import threading +import time +from typing import Any, Dict, List, Optional, Tuple + +import duckdb + +_instance: Optional["QueryExecutor"] = None +_lock = threading.Lock() + + +def _strip_terminators(query: str) -> str: + """Trim surrounding whitespace and any trailing semicolons. + + Trailing ``;`` breaks queries that are wrapped in a subquery + (``SELECT COUNT(*) FROM ({query}) t``) for correctness checksums. + """ + return query.strip().rstrip(";").strip() + + +class QueryExecutor: + """ + Runs SQL against an in-memory DuckDB database with realistic + synthetic data. Provides execution timing, result correctness + checks, and EXPLAIN plans ā€” all used by the reward function. + """ + + def __init__(self) -> None: + # Build the tables once on a writable connection to a temp file, then + # reopen the database read-only. Every query ā€” trusted task SQL and + # agent rewrites alike ā€” runs against the read-only connection, so the + # DuckDB engine itself rejects any write (no keyword heuristics, and DML + # inside a CTE cannot slip through on any DuckDB version). A lock + # serializes access because a single DuckDB connection is not safe for + # concurrent use by the server's worker pool. + self._dir = tempfile.mkdtemp(prefix="sql_optim_") + self._path = os.path.join(self._dir, "sql_optim.duckdb") + builder = duckdb.connect(self._path, config={"threads": "2"}) + self._build_tables(builder) + builder.close() + + self.conn = duckdb.connect(self._path, read_only=True, config={"threads": "2"}) + self._exec_lock = threading.Lock() + atexit.register(self._cleanup) + + def _cleanup(self) -> None: + """Close the connection and remove the temp database directory.""" + try: + self.conn.close() + except Exception: + pass + shutil.rmtree(self._dir, ignore_errors=True) + + # ā”€ā”€ Schema Setup ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + def _build_tables(self, conn: "duckdb.DuckDBPyConnection") -> None: + """Create and populate all four synthetic tables on *conn*.""" + + # users ā€” 10k rows + conn.execute(""" + CREATE TABLE users AS + SELECT + i AS id, + 'u' || i || '@mail.com' AS email, + CASE i % 3 + WHEN 0 THEN 'premium' + WHEN 1 THEN 'free' + ELSE 'enterprise' END AS tier, + CASE i % 5 + WHEN 0 THEN 'US' WHEN 1 THEN 'EU' + WHEN 2 THEN 'IN' WHEN 3 THEN 'UK' + ELSE 'AU' END AS region, + CASE i % 2 WHEN 0 THEN 'premium' ELSE 'basic' END AS plan, + DATE '2020-01-01' + CAST(i AS INTEGER) AS created_at + FROM generate_series(1, 10000) t(i) + """) + + # orders ā€” 500k rows + conn.execute(""" + CREATE TABLE orders AS + SELECT + i AS id, + 1 + (i % 10000) AS customer_id, + (i % 100) + 1 AS product_id, + CASE i % 4 + WHEN 0 THEN 'completed' WHEN 1 THEN 'pending' + WHEN 2 THEN 'cancelled' ELSE 'shipped' END AS status, + ROUND((i % 1000) * 1.5 + 49.99, 2) AS total, + DATE '2023-01-01' + CAST(i % 730 AS INTEGER) AS created_at + FROM generate_series(1, 500000) t(i) + """) + + # products ā€” 1k rows + conn.execute(""" + CREATE TABLE products AS + SELECT + i AS id, + 'Product_' || i AS name, + CASE i % 5 + WHEN 0 THEN 'Electronics' WHEN 1 THEN 'Clothing' + WHEN 2 THEN 'Food' WHEN 3 THEN 'Books' + ELSE 'Sports' END AS category, + ROUND((i % 500) + 9.99, 2) AS price + FROM generate_series(1, 1000) t(i) + """) + + # events ā€” 1M rows + conn.execute(""" + CREATE TABLE events AS + SELECT + i AS id, + 1 + (i % 10000) AS user_id, + 'sess_' || (i % 50000) AS session_id, + CASE i % 6 + WHEN 0 THEN 'purchase' WHEN 1 THEN 'view' + WHEN 2 THEN 'click' WHEN 3 THEN 'signup' + WHEN 4 THEN 'logout' ELSE 'search' END AS event_type, + DATE '2024-01-01' + CAST(i % 365 AS INTEGER) AS occurred_at + FROM generate_series(1, 1000000) t(i) + """) + + # ā”€ā”€ Execution helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + def _run( + self, query: str, runs: int = 3 + ) -> Tuple[float, Optional[List], Optional[str]]: + """ + Execute *query* up to *runs* times on the read-only connection. + Returns (median_ms, rows, error_or_None). A write raises and is returned + as the error string. + """ + timings: List[float] = [] + rows: Optional[List] = None + + with self._exec_lock: + for _ in range(runs): + try: + t0 = time.perf_counter() + rows = self.conn.execute(query).fetchall() + timings.append((time.perf_counter() - t0) * 1000.0) + except Exception as exc: + return 99_999.0, None, str(exc) + + timings.sort() + return round(timings[len(timings) // 2], 3), rows, None + + def _checksum( + self, query: str + ) -> Tuple[Optional[int], Optional[int], Optional[str]]: + """ + Compute a deterministic (row-order-independent) checksum. + Returns (row_count, checksum, error). + + BIT_XOR is commutative+associative ā€” order-independent fingerprint. + Falls back to count-only if the DuckDB version doesn't support the function. + """ + # Try BIT_XOR of a numeric hash (portable across DuckDB versions) + for sql_template in [ + # Option 1: BIT_XOR of md5 prefix cast to integer + ( + "SELECT COUNT(*) AS cnt, " + "BIT_XOR(CAST(('0x' || LEFT(md5(CAST(t AS VARCHAR)), 15)) AS UBIGINT)) AS chk " + "FROM ({query}) t" + ), + # Option 2: sum of hash (order-independent since sum is commutative) + ( + "SELECT COUNT(*) AS cnt, " + "SUM(hash(CAST(t AS VARCHAR)) % 9999999999) AS chk " + "FROM ({query}) t" + ), + ]: + try: + wrapped = sql_template.format(query=query) + with self._exec_lock: + result = self.conn.execute(wrapped).fetchone() + return result[0], result[1], None + except Exception: + continue + # Final fallback: count only + try: + with self._exec_lock: + cnt = self.conn.execute(f"SELECT COUNT(*) FROM ({query}) t").fetchone()[ + 0 + ] + return cnt, None, None + except Exception as exc: + return None, None, str(exc) + + # ā”€ā”€ Public API ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + def compare(self, original: str, optimized: str) -> Dict[str, Any]: + """ + Execute both queries, measure real timing, check correctness. + + Returns a dict with: + original_ms, optimized_ms, speedup, + results_match, original_rows, optimized_rows, + original_error, optimized_error, verdict + """ + # Normalize (drop trailing ``;``) so both direct execution and the + # subquery-wrapped checksum are valid. + original = _strip_terminators(original) + optimized = _strip_terminators(optimized) + + orig_ms, orig_rows, orig_err = self._run(original) + + # The optimized query is agent-authored, but the read-only connection is + # the guarantee: any write (DDL/DML, in a CTE, or after a `;`) is rejected + # by the DuckDB engine and surfaces as `optimized_error`. No structural + # pre-check is needed, which also avoids false rejections of valid + # rewrites (leading comments, a `;` inside a string literal, etc.). + opt_ms, opt_rows, opt_err = self._run(optimized) + + # ā”€ā”€ Correctness: do both queries return the same data? ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # Use a DuckDB-level checksum (order-independent) to avoid + # false negatives from non-deterministic row ordering in parallel + # window function queries on large tables. + results_match = False + if orig_rows is not None and opt_rows is not None: + try: + if len(orig_rows) != len(opt_rows): + results_match = False + elif len(orig_rows) == 0: + results_match = True + elif len(orig_rows) <= 50_000: + # Small/medium: full sorted comparison (precise) + orig_s = sorted(str(r) for r in orig_rows) + opt_s = sorted(str(r) for r in opt_rows) + results_match = orig_s == opt_s + else: + # Large result sets: use SQL-level hash checksum + # (deterministic regardless of row ordering / thread count) + o_cnt, o_chk, o_err2 = self._checksum(original) + p_cnt, p_chk, p_err2 = self._checksum(optimized) + if o_err2 or p_err2: + # Checksum failed ā€” fall back to row count + results_match = len(orig_rows) == len(opt_rows) + else: + results_match = (o_cnt == p_cnt) and (o_chk == p_chk) + except Exception: + results_match = len(orig_rows) == len(opt_rows) + + # ā”€ā”€ Speedup ratio ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + speedup = 1.0 + if opt_ms > 0 and orig_ms < 90_000: + speedup = round(orig_ms / opt_ms, 3) + + # ā”€ā”€ Human-readable verdict ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + if opt_err: + verdict = f"[FAIL] Optimized query error: {opt_err[:120]}" + elif results_match and speedup >= 2.0: + verdict = f"[OK] {speedup:.1f}x faster with correct results" + elif results_match and speedup >= 1.0: + verdict = ( + f"[WARN] Correct results but only {speedup:.1f}x speedup -- dig deeper" + ) + elif not results_match and speedup >= 2.0: + verdict = ( + f"[WARN] {speedup:.1f}x faster but results don't match -- fix the logic" + ) + else: + verdict = f"[FAIL] {speedup:.1f}x -- no meaningful improvement" + + return { + "original_ms": orig_ms, + "optimized_ms": opt_ms, + "speedup": speedup, + "results_match": results_match, + "original_rows": len(orig_rows) if orig_rows is not None else 0, + "optimized_rows": len(opt_rows) if opt_rows is not None else 0, + "original_error": orig_err, + "optimized_error": opt_err, + "verdict": verdict, + } + + def explain(self, query: str) -> str: + """Return EXPLAIN output for a query.""" + try: + with self._exec_lock: + rows = self.conn.execute( + f"EXPLAIN {_strip_terminators(query)}" + ).fetchall() + return "\n".join(str(r[1]) for r in rows) + except Exception as exc: + return f"EXPLAIN error: {exc}" + + @property + def table_stats(self) -> Dict[str, int]: + tables = ["users", "orders", "products", "events"] + with self._exec_lock: + return { + t: self.conn.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0] + for t in tables + } + + +# ā”€ā”€ Singleton accessor ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def get_executor() -> QueryExecutor: + """Return the process-level singleton (lazy init, thread-safe).""" + global _instance + if _instance is None: + with _lock: + if _instance is None: + _instance = QueryExecutor() + return _instance diff --git a/envs/sql_optim_env/server/graders.py b/envs/sql_optim_env/server/graders.py new file mode 100644 index 000000000..f44b77ec7 --- /dev/null +++ b/envs/sql_optim_env/server/graders.py @@ -0,0 +1,219 @@ +""" +graders.py ā€” Execution-Grounded Reward Function +================================================= +What makes this environment unique: reward is computed from REAL +DuckDB execution results, not just keyword heuristics. + +Scoring breakdown (sums to 1.0): + Real Execution Speedup 35% ā€” actual timing ratio from DuckDB + Result Correctness 20% ā€” both queries return identical data? + Issue Detection 25% ā€” keyword match vs ground truth + Approval Correctness 8% ā€” correctly flags query as bad? + Summary Quality 7% ā€” is the written analysis thorough? + Severity Labels 5% ā€” are severity values present? + +Optional ``GradeMask`` (keyword arg ``mask=``) zeroes components for ablations; +production calls omit ``mask`` (full scoring, including the 0.02 minimum when +appropriate). +""" + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +try: + # In-repo imports (running from the OpenEnv repository) + from ..models import SQLOptimAction as Action + from .executor import get_executor + from .scoring_models import Reward +except ImportError as e: + if "relative import" not in str(e) and "no known parent package" not in str(e): + raise + # Standalone imports (running via uvicorn server.app:app) + from models import SQLOptimAction as Action + from server.executor import get_executor + from server.scoring_models import Reward + + +@dataclass(frozen=True) +class GradeMask: + """Toggle reward components (for ablations). All True = production grading.""" + + execution_speedup: bool = True + result_correctness: bool = True + issue_detection: bool = True + approval_correctness: bool = True + summary_quality: bool = True + severity_labels: bool = True + + +# ā”€ā”€ Helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _kw_match(text: str, keywords: List[str]) -> bool: + t = text.lower() + return any(kw.lower() in t for kw in keywords) + + +def _suggestions_text(action: Action) -> str: + # Only the agent's analysis (summary + structured suggestions) counts toward + # issue detection ā€” deliberately NOT `optimized_query`. Otherwise echoing the + # slow SQL (which contains the anti-pattern tokens) would farm keyword credit + # without any real analysis. + parts = [action.summary, action.estimated_improvement] + for s in action.suggestions: + parts += [ + str(s.get("issue_type", "")), + str(s.get("description", "")), + str(s.get("fix", "")), + str(s.get("severity", "")), + ] + return " ".join(parts) + + +# ā”€ā”€ Speedup ā†’ score mapping ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _speedup_score(speedup: float, has_error: bool) -> float: + """Map real speedup ratio to a score in [0.0, 0.35].""" + if has_error: + return 0.0 + if speedup >= 15.0: + return 0.35 + if speedup >= 8.0: + return 0.30 + if speedup >= 4.0: + return 0.25 + if speedup >= 2.0: + return 0.18 + if speedup >= 1.2: + return 0.10 + if speedup >= 0.9: # slightly slower ā€” acceptable + return 0.04 + return 0.0 # regression + + +# ā”€ā”€ Main grader ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def grade( + task_data: Dict[str, Any], + action: Action, + *, + mask: Optional[GradeMask] = None, +) -> Reward: + original_query: str = task_data["sql_query"] + optimized_query: str = (action.optimized_query or "").strip() + ground_truth: List[Dict[str, Any]] = task_data["ground_truth_issues"] + full_text = _suggestions_text(action) + + # ā”€ā”€ 1. Real Execution (0.0ā€“0.35) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + exec_info: Dict[str, Any] = {} + speedup_sc = 0.0 + correctness_sc = 0.0 + exec_feedback: List[str] = [] + + if optimized_query: + try: + ex = get_executor() + exec_info = ex.compare(original_query, optimized_query) + speedup = exec_info.get("speedup", 1.0) + r_match = exec_info.get("results_match", False) + opt_err = exec_info.get("optimized_error") + + # 1a. Speedup score ā€” only credited when the rewrite is correct, so a + # trivially fast but wrong query (e.g. `SELECT 1` against a heavy + # original) earns nothing for speed. + speedup_sc = _speedup_score(speedup, has_error=bool(opt_err) or not r_match) + + # 1b. Correctness score (0.0-0.20) + if opt_err: + correctness_sc = 0.0 + elif r_match: + correctness_sc = 0.20 + elif exec_info.get("optimized_rows", 0) > 0: + # Query ran but different results -- partial credit + correctness_sc = 0.05 + + # Feedback lines + exec_feedback = [ + "\n[DuckDB Execution Results]", + f" Original : {exec_info['original_ms']:.1f} ms " + f"({exec_info['original_rows']} rows)", + f" Optimized : {exec_info['optimized_ms']:.1f} ms " + f"({exec_info['optimized_rows']} rows)", + f" Speedup : {speedup:.2f}x", + f" Correct? : {'YES' if r_match else 'NO -- results differ'}", + f" Verdict : {exec_info.get('verdict', '')}", + ] + if opt_err: + exec_feedback.append(f" SQL Error : {opt_err[:200]}") + + except Exception as exc: + exec_feedback = [f"\n[WARN] Execution engine error: {exc}"] + + # ā”€ā”€ 2. Issue Detection (0.0ā€“0.25) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + detected = 0 + detection_fb: List[str] = ["\n[Issue Detection]"] + for gt in ground_truth: + found = _kw_match(full_text, gt["keywords"]) + if found: + detected += 1 + detection_fb.append(f" [FOUND] {gt['type']} (line ~{gt['line']})") + else: + detection_fb.append(f" [MISS ] {gt['type']} (line ~{gt['line']})") + detection_sc = (detected / len(ground_truth)) * 0.25 if ground_truth else 0.0 + + # ā”€ā”€ 3. Approval Correctness (0.0ā€“0.08) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + expected_approved = task_data.get("approved_expected", False) + approval_sc = 0.08 if action.approved == expected_approved else 0.0 + + # ā”€ā”€ 4. Summary Quality (0.0ā€“0.07) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + summary_sc = 0.0 + slen = len(action.summary) + if slen > 50: + summary_sc = 0.03 + if slen > 120: + summary_sc = 0.07 + + # ā”€ā”€ 5. Severity Labels (0.0ā€“0.05) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + sev_kw = ["critical", "high", "medium", "low"] + has_sev = any( + _kw_match(str(s.get("severity", "")), sev_kw) for s in action.suggestions + ) + severity_sc = 0.05 if has_sev else 0.0 + + breakdown = { + "execution_speedup": round(speedup_sc, 4), + "result_correctness": round(correctness_sc, 4), + "issue_detection": round(detection_sc, 4), + "approval_correctness": round(approval_sc, 4), + "summary_quality": round(summary_sc, 4), + "severity_labels": round(severity_sc, 4), + } + + # ā”€ā”€ Total (optional component mask for ablations) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + m = mask or GradeMask() + contrib = {k: (v if getattr(m, k) else 0.0) for k, v in breakdown.items()} + total = round(min(max(sum(contrib.values()), 0.0), 1.0), 4) + if mask is None and total == 0.0 and action.suggestions: + total = 0.02 # minimum signal for any submission (production only) + + feedback = "\n".join( + exec_feedback + + detection_fb + + [ + f"\n Suggestions submitted: {len(action.suggestions)} " + f"(expected ~{len(ground_truth)})", + f" Approval: {'āœ…' if action.approved == expected_approved else 'āŒ'} " + f"(got {'approved' if action.approved else 'rejected'}, " + f"expected {'approved' if expected_approved else 'rejected'})", + f"\nšŸ† Total score: {total:.4f}", + ] + ) + + return Reward( + score=total, + breakdown=breakdown, + feedback=feedback, + execution=exec_info or None, + ) diff --git a/envs/sql_optim_env/server/scoring_models.py b/envs/sql_optim_env/server/scoring_models.py new file mode 100644 index 000000000..c4b4c6310 --- /dev/null +++ b/envs/sql_optim_env/server/scoring_models.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Server-internal scoring types for the SQL Query Optimization environment. + +`Reward` is produced by [`grade`][envs.sql_optim_env.server.graders.grade] and +consumed by the environment, which folds it into the observation's `reward` +field and `metadata`. It is not a wire type and never leaves the server. +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + +from pydantic import BaseModel, Field + + +class Reward(BaseModel): + """ + Composite, execution-grounded reward for one optimization step. + + Attributes: + score (`float`): + Composite reward in `[0.0, 1.0]`. + breakdown (`dict[str, float]`): + Per-criterion contribution to `score` (execution speedup, result + correctness, issue detection, approval, summary, severity). + feedback (`str`): + Human-readable explanation, including real DuckDB timings. + execution (`dict`, *optional*): + The DuckDB `compare()` result this reward was scored from (timings, + speedup, correctness), or `None` when no query was executed. Carried + here so the environment can surface it without re-running the queries. + """ + + score: float = Field(..., ge=0.0, le=1.0) + breakdown: Dict[str, float] = Field(default_factory=dict) + feedback: str = "" + execution: Optional[Dict[str, Any]] = None diff --git a/envs/sql_optim_env/server/sql_optim_environment.py b/envs/sql_optim_env/server/sql_optim_environment.py new file mode 100644 index 000000000..5006a111a --- /dev/null +++ b/envs/sql_optim_env/server/sql_optim_environment.py @@ -0,0 +1,194 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Server-side environment for SQL query optimization. + +The agent receives a slow SQL query plus schema context and returns an +optimized rewrite. The environment executes both the original and optimized +queries against a real in-memory DuckDB database, measures the actual speedup +and result-correctness, and grades the result ā€” reward lives entirely inside +the environment, per OpenEnv's design. +""" + +from __future__ import annotations + +import uuid +from typing import Any, Optional + +from openenv.core.env_server import Environment + +# Support both in-repo and standalone imports +try: + # In-repo imports (running from the OpenEnv repository) + from ..models import SQLOptimAction, SQLOptimObservation, SQLOptimState + from .graders import grade + from .tasks import TASKS +except ImportError as e: + if "relative import" not in str(e) and "no known parent package" not in str(e): + raise + # Standalone imports (running via uvicorn server.app:app) + from models import SQLOptimAction, SQLOptimObservation, SQLOptimState + from server.graders import grade + from server.tasks import TASKS + +DEFAULT_TASK_ID = "task_1_basic_antipatterns" + + +class SQLOptimEnvironment(Environment): + """ + Execution-grounded SQL optimization environment. + + Multi-step episodes: `issues_found_so_far` accumulates the issue types the + agent has flagged, and `last_execution` carries the real DuckDB timing + comparison back into the observation so the agent can refine its rewrite on + later steps. An episode ends after the task's `max_steps` or once a step + scores `>= 0.95`. + """ + + def __init__(self) -> None: + super().__init__() + self._task_data: Optional[dict] = None + self._step_count: int = 0 + self._done: bool = False + self._cumulative_reward: float = 0.0 + self._issues_found: list[str] = [] + self._last_execution: Optional[dict] = None + self._episode_id: str = "" + self._state: SQLOptimState = SQLOptimState() + self.reset() + + def reset( + self, + seed: Optional[int] = None, + episode_id: Optional[str] = None, + task_id: Optional[str] = None, + **kwargs: Any, + ) -> SQLOptimObservation: + """ + Start a new episode on the requested task. + + Args: + seed (`int`, *optional*): + Accepted for Gym API compatibility; the task data is + deterministic, so it does not affect the episode. + episode_id (`str`, *optional*): + Custom episode identifier. A UUID is generated when omitted. + task_id (`str`, *optional*, defaults to `"task_1_basic_antipatterns"`): + Which task to load. Must be a key of `TASKS`. + + Returns: + `SQLOptimObservation`: The initial observation for the task. + """ + task_id = task_id or DEFAULT_TASK_ID + if task_id not in TASKS: + raise ValueError( + f"Unknown task_id '{task_id}'. Valid: {list(TASKS.keys())}" + ) + self._task_data = TASKS[task_id] + self._step_count = 0 + self._done = False + self._cumulative_reward = 0.0 + self._issues_found = [] + self._last_execution = None + self._episode_id = episode_id or str(uuid.uuid4()) + self._sync_state() + return self._make_obs() + + def step(self, action: SQLOptimAction) -> SQLOptimObservation: + """ + Grade one optimization attempt and advance the episode. + + Args: + action (`SQLOptimAction`): + The agent's suggestions and rewritten `optimized_query`. + + Returns: + `SQLOptimObservation`: The next observation, with the step reward on + `reward`, the composite breakdown and execution comparison under + `metadata`, and `done` set when the episode terminates. + """ + if self._task_data is None: + raise RuntimeError("No active episode - call reset() first.") + if self._done: + raise RuntimeError("Episode finished - call reset() to start a new one.") + + self._step_count += 1 + + # Grade (runs DuckDB internally). Reuse the comparison it already ran + # rather than executing both queries a second time. `execution` is + # `None` when the rewrite was empty or failed, so a skipped step never + # reports the previous step's stale comparison. + reward = grade(self._task_data, action) + self._cumulative_reward += reward.score + self._last_execution = reward.execution + + # Track flagged issue types for progressive context. + for suggestion in action.suggestions: + itype = suggestion.get("issue_type", "") + if itype and itype not in self._issues_found: + self._issues_found.append(itype) + + max_steps = self._task_data["max_steps"] + self._done = self._step_count >= max_steps or reward.score >= 0.95 + self._sync_state() + + return self._make_obs( + reward=reward.score, + done=self._done, + metadata={ + "reward_breakdown": reward.breakdown, + "feedback": reward.feedback, + "execution": self._last_execution, + "cumulative_reward": round(self._cumulative_reward, 4), + "issues_found": len(self._issues_found), + }, + ) + + @property + def state(self) -> SQLOptimState: + """Current [`SQLOptimState`] for this episode.""" + return self._state + + # ā”€ā”€ Internal ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + def _make_obs( + self, + reward: Optional[float] = None, + done: bool = False, + metadata: Optional[dict] = None, + ) -> SQLOptimObservation: + d = self._task_data + return SQLOptimObservation( + task_id=d["task_id"], + task_name=d["task_name"], + task_description=d["task_description"], + sql_query=d["sql_query"], + schema_info=d["schema_info"], + dialect=d.get("dialect", "duckdb/postgresql"), + difficulty=d["difficulty"], + step_count=self._step_count, + max_steps=d["max_steps"], + issues_found_so_far=list(self._issues_found), + last_execution=self._last_execution, + reward=reward, + done=done, + metadata=metadata or {}, + ) + + def _sync_state(self) -> None: + if self._task_data is None: + self._state = SQLOptimState(episode_id=self._episode_id) + return + self._state = SQLOptimState( + episode_id=self._episode_id, + step_count=self._step_count, + task_id=self._task_data["task_id"], + max_steps=self._task_data["max_steps"], + episode_done=self._done, + cumulative_reward=round(self._cumulative_reward, 4), + current_task=self._task_data["task_name"], + ) diff --git a/envs/sql_optim_env/server/tasks.py b/envs/sql_optim_env/server/tasks.py new file mode 100644 index 000000000..5616d3384 --- /dev/null +++ b/envs/sql_optim_env/server/tasks.py @@ -0,0 +1,523 @@ +""" +tasks.py ā€” SQL Query Optimization Tasks +======================================== +Five tasks of increasing difficulty, each with a DuckDB-executable +"bad" query (stored in sql_query) that agents must analyze and rewrite. + +All queries run against the executor's synthetic tables: + users (10,000 rows) ā€” id, email, tier, region, plan, created_at + orders (500,000 rows) ā€” id, customer_id, product_id, status, total, created_at + products (1,000 rows) ā€” id, name, category, price + events (1,000,000 rows) ā€” id, user_id, session_id, event_type, occurred_at +""" + +from typing import Any, Dict + +TASKS: Dict[str, Dict[str, Any]] = { + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # TASK 1 ā€” EASY: Basic Anti-pattern Detection + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + "task_1_basic_antipatterns": { + "task_id": "task_1_basic_antipatterns", + "task_name": "Basic SQL Anti-pattern Detection", + "task_description": ( + "Analyze the SQL query below for common anti-patterns that destroy performance. " + "Your rewrite must return the SAME rows and columns: replace the non-sargable " + "predicates ā€” `CAST(customer_id AS VARCHAR) = '5000'` with `customer_id = 5000`, " + "and `year(created_at) = 2024` with a `created_at >= '2024-01-01' AND " + "created_at < '2025-01-01'` range ā€” so DuckDB can prune. " + "Separately identify ā€” but do NOT remove, since it changes the output ā€” the " + "SELECT * that fetches every column from 500k rows. " + "For each issue report: issue_type, line, description, severity, and a concrete fix. " + "Your optimized_query is EXECUTED against real data; speedup is credited only " + "when the results still match." + ), + "difficulty": "easy", + "dialect": "duckdb/postgresql", + "max_steps": 3, + "schema_info": ( + "Table: orders (500,000 rows)\n" + " id INT, customer_id INT, product_id INT,\n" + " status VARCHAR, total DECIMAL, created_at DATE\n\n" + "No indexes defined (DuckDB uses columnar min/max pruning when columns " + "are not wrapped in functions).\n" + "Scan cost: ~500k rows Ɨ all columns with SELECT *" + ), + "sql_query": ( + "SELECT *\n" + "FROM orders\n" + "WHERE CAST(customer_id AS VARCHAR) = '5000'\n" + " AND year(created_at) = 2024;" + ), + "ground_truth_issues": [ + { + "type": "select_star", + "line": 1, + "keywords": [ + "select *", + "star", + "all columns", + "unnecessary columns", + "column projection", + "specify columns", + "bandwidth", + ], + }, + { + "type": "non_sargable_cast", + "line": 3, + "keywords": [ + "cast", + "varchar", + "type cast", + "type conversion", + "non-sargable", + "sargable", + "integer comparison", + "string comparison", + "prevents", + "pruning", + ], + }, + { + "type": "function_on_date_column", + "line": 4, + "keywords": [ + "year(", + "function on column", + "non-sargable", + "date range", + "between", + "extract", + "full scan", + "date filter", + ], + }, + ], + "approved_expected": False, + }, + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # TASK 2 ā€” MEDIUM: N+1 Correlated Subqueries + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + "task_2_correlated_subqueries": { + "task_id": "task_2_correlated_subqueries", + "task_name": "N+1 Correlated Subquery Elimination", + "task_description": ( + "The query below uses three correlated scalar subqueries ā€” each one scans " + "the entire orders table (500k rows) once per premium user (~3,300 users). " + "That's ~10 million row reads just for aggregation. " + "Identify the N+1 pattern, explain why each subquery is harmful, " + "and rewrite the query as a single aggregation JOIN. " + "Your optimized_query will be executed; results must match the original." + ), + "difficulty": "medium", + "dialect": "duckdb/postgresql", + "max_steps": 4, + "schema_info": ( + "Table: users (10,000 rows)\n" + " id INT, email VARCHAR, tier VARCHAR, region VARCHAR,\n" + " plan VARCHAR, created_at DATE\n\n" + "Table: orders (500,000 rows)\n" + " id INT, customer_id INT, product_id INT,\n" + " status VARCHAR, total DECIMAL, created_at DATE\n\n" + "Premium users: ~3,300 | Orders per user avg: 50\n" + "Worst-case scans: 3 subqueries Ɨ 3,300 users Ɨ 500k rows = ~5B row reads" + ), + "sql_query": ( + "SELECT\n" + " u.email,\n" + " u.region,\n" + " (SELECT COUNT(*)\n" + " FROM orders o\n" + " WHERE o.customer_id = u.id AND o.status = 'completed') AS completed_orders,\n" + " (SELECT SUM(o.total)\n" + " FROM orders o\n" + " WHERE o.customer_id = u.id\n" + " AND o.created_at >= DATE '2024-01-01') AS ytd_spend,\n" + " (SELECT total\n" + " FROM orders o\n" + " WHERE o.customer_id = u.id\n" + " ORDER BY created_at DESC, id DESC LIMIT 1) AS last_order_amount\n" + "FROM users u\n" + "WHERE u.tier = 'premium';" + ), + "ground_truth_issues": [ + { + "type": "correlated_subquery_count", + "line": 4, + "keywords": [ + "correlated", + "subquery", + "per row", + "n+1", + "each user", + "repeated scan", + "join", + "aggregation", + ], + }, + { + "type": "correlated_subquery_sum", + "line": 7, + "keywords": [ + "correlated", + "subquery", + "per row", + "n+1", + "each user", + "repeated scan", + "join", + "group by", + ], + }, + { + "type": "correlated_subquery_limit", + "line": 11, + "keywords": [ + "correlated", + "subquery", + "limit 1", + "order by", + "lateral", + "row_number", + "rank", + "window function", + "per row", + ], + }, + { + "type": "missing_aggregation_join", + "line": 16, + "keywords": [ + "left join", + "group by", + "aggreg", + "single pass", + "coalesce", + "join aggregat", + ], + }, + ], + "approved_expected": False, + }, + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # TASK 3 ā€” MEDIUM-HARD: Wildcard LIKE Full Scan + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + "task_3_wildcard_scan": { + "task_id": "task_3_wildcard_scan", + "task_name": "Wildcard LIKE & Projection Optimization", + "task_description": ( + "The query scans all 1,000,000 events rows with leading-wildcard LIKE " + "patterns that disable columnar zone-map pruning, plus a redundant OR branch " + "that never matches any row. " + "Your rewrite must return the SAME rows and columns: replace the " + "leading-wildcard filter with exact equality (`event_type = 'purchase'`), " + "which is provably equivalent here and lets DuckDB prune. " + "Separately identify ā€” but do NOT apply, since they would change the output ā€” " + "the SELECT * on a million-row table and the derived columns computed for " + "every row." + ), + "difficulty": "medium-hard", + "dialect": "duckdb/postgresql", + "max_steps": 4, + "schema_info": ( + "Table: events (1,000,000 rows)\n" + " id INT, user_id INT, session_id VARCHAR,\n" + " event_type VARCHAR, occurred_at DATE\n\n" + "Distinct event_type values: purchase, view, click, signup, logout, search\n" + "Only 'purchase' contains the substring 'purchase' and no event_type contains\n" + "'buy', so the filter is equivalent to `event_type = 'purchase'`.\n" + "Leading-wildcard LIKE on 1M rows forces a full scan; exact equality enables\n" + "columnar zone-map pruning." + ), + "sql_query": ( + "SELECT\n" + " *,\n" + " CAST(id AS VARCHAR) || '_' || event_type AS event_key,\n" + " upper(event_type) AS event_type_upper\n" + "FROM events\n" + "WHERE event_type LIKE '%purchase%'\n" + " OR event_type LIKE '%buy%';" + ), + "ground_truth_issues": [ + { + "type": "leading_wildcard_like", + "line": 6, + "keywords": [ + "leading wildcard", + "like '%", + "full scan", + "exact match", + "equality", + "pruning disabled", + "wildcard", + "zone map", + ], + }, + { + "type": "or_expands_to_full_scan", + "line": 7, + "keywords": [ + "or condition", + "union", + "separate queries", + "or expands", + "full scan", + "like '%buy%'", + "redundant", + ], + }, + { + "type": "select_star_large_table", + "line": 2, + "keywords": [ + "select *", + "1 million", + "all columns", + "projection", + "column pruning", + "unnecessary", + "bandwidth", + ], + }, + { + "type": "pre_filter_computed_columns", + "line": 3, + "keywords": [ + "computed column", + "derived", + "upper(", + "cast(", + "concatenat", + "before filter", + "pre-filter", + "push down", + "CTE", + ], + }, + ], + "approved_expected": False, + }, + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # TASK 4 ā€” HARD: Implicit Cross Join + Repeated Scalar Subqueries + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + "task_4_implicit_join": { + "task_id": "task_4_implicit_join", + "task_name": "Implicit Cross Join & Scalar Subquery Elimination", + "task_description": ( + "This query uses comma-separated FROM (implicit cross join syntax) and " + "two uncorrelated scalar subqueries that compute global constants over the " + "orders table (neither references the outer query). " + "Identify: implicit cross join risk (comma in FROM clause), " + "two uncorrelated scalar subqueries computing global stats that are clearer " + "hoisted into a CTE, " + "and the GROUP BY without an explicit JOIN. " + "Rewrite using explicit INNER JOIN and a CTE for the global stats " + "so the intent is explicit and they are computed once." + ), + "difficulty": "hard", + "dialect": "duckdb/postgresql", + "max_steps": 5, + "schema_info": ( + "Table: users (10,000 rows) ā€” id, email, tier, region, plan, created_at\n" + "Table: orders (500,000 rows) ā€” id, customer_id, product_id, status, total, created_at\n\n" + "Join: users.id = orders.customer_id\n" + "Implicit join (comma syntax) risk: if WHERE predicate is missing,\n" + "produces a Cartesian product of 10k Ɨ 500k = 5 BILLION rows.\n" + "Scalar subqueries: uncorrelated global aggregates over all 500k orders;\n" + "clearer and guaranteed single-eval when hoisted into a CTE." + ), + "sql_query": ( + "SELECT\n" + " u.region,\n" + " u.plan,\n" + " COUNT(*) AS total_orders,\n" + " SUM(o.total) AS revenue,\n" + " (SELECT AVG(total) FROM orders) AS global_avg,\n" + " (SELECT MAX(total) FROM orders WHERE status = 'completed') AS max_deal\n" + "FROM users u, orders o\n" + "WHERE u.id = o.customer_id\n" + " AND o.status IN ('completed', 'shipped')\n" + "GROUP BY u.region, u.plan;" + ), + "ground_truth_issues": [ + { + "type": "implicit_cross_join", + "line": 8, + "keywords": [ + "implicit", + "cross join", + "comma join", + "explicit join", + "inner join", + "cartesian", + "comma in from", + ], + }, + { + "type": "uncorrelated_scalar_subquery_avg", + "line": 6, + "keywords": [ + "scalar subquery", + "uncorrelated", + "global constant", + "hoist", + "cte", + "with clause", + "pre-compute", + "global avg", + ], + }, + { + "type": "uncorrelated_scalar_subquery_max", + "line": 7, + "keywords": [ + "scalar subquery", + "uncorrelated", + "global constant", + "max deal", + "cte", + "pre-compute", + "compute once", + "constant", + ], + }, + { + "type": "missing_explicit_join", + "line": 8, + "keywords": [ + "inner join", + "explicit", + "on clause", + "join condition", + "readable", + "maintainable", + "ansi sql", + ], + }, + ], + "approved_expected": False, + }, + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # TASK 5 ā€” EXPERT: Window Function Over Entire 1M-Row Table + # ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + "task_5_window_functions": { + "task_id": "task_5_window_functions", + "task_name": "Window Function & Full-Scan Audit", + "task_description": ( + "Five window functions are computed over ALL 1,000,000 events rows. Three of " + "them share `PARTITION BY user_id` and can be consolidated so that partition is " + "built once instead of repeatedly. " + "Your rewrite must return identical rows and columns: keep all output columns " + "but consolidate the window functions that share `PARTITION BY user_id` " + "(result-preserving). " + "Separately identify ā€” but do NOT apply, since they change the result ā€” the " + "absence of pre-filtering (a WHERE would drop rows) and the global " + "RANK() OVER (ORDER BY occurred_at) that sorts the entire table." + ), + "difficulty": "expert", + "dialect": "duckdb/postgresql", + "max_steps": 5, + "schema_info": ( + "Table: events (1,000,000 rows)\n" + " id INT, user_id INT, session_id VARCHAR,\n" + " event_type VARCHAR, occurred_at DATE\n\n" + "Window function cost: each OVER() = full sort/hash pass over 1M rows\n" + "5 window functions = 5 full passes before any filtering\n" + "Global RANK(): sorts all 1M rows globally ā€” most expensive operation\n" + "Filtering to 'purchase' events first reduces dataset to ~167k rows (1/6)" + ), + "sql_query": ( + "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, id) AS recency_rank,\n" + " RANK() OVER (ORDER BY occurred_at DESC, id) AS global_rank,\n" + " SUM(CASE WHEN event_type = 'purchase' THEN 1 ELSE 0 END)\n" + " OVER (PARTITION BY user_id) AS user_purchases\n" + "FROM events;" + ), + "ground_truth_issues": [ + { + "type": "no_pre_filter", + "line": 11, + "keywords": [ + "no where", + "no filter", + "full table", + "1 million", + "all rows", + "pre-filter", + "filter first", + "cte", + "with clause", + ], + }, + { + "type": "global_rank_no_partition", + "line": 8, + "keywords": [ + "rank() over", + "global rank", + "no partition", + "entire table", + "full sort", + "expensive", + "global ordering", + "remove", + ], + }, + { + "type": "redundant_window_functions", + "line": 5, + "keywords": [ + "5 window", + "multiple over", + "redundant", + "merge", + "combine", + "single pass", + "same partition", + "consolidate", + ], + }, + { + "type": "count_vs_conditional_sum", + "line": 9, + "keywords": [ + "case when", + "sum case", + "count filter", + "filter clause", + "count(*) filter", + "simpler", + "merge with", + ], + }, + ], + "approved_expected": False, + }, +} + + +def get_task_list(): + return [ + { + "task_id": t["task_id"], + "task_name": t["task_name"], + "difficulty": t["difficulty"], + "max_steps": t["max_steps"], + "description": t["task_description"], + "action_schema": { + "suggestions": "List of {issue_type, line, description, severity, fix}", + "optimized_query": "str ā€” complete rewritten SQL (will be EXECUTED for real timing)", + "summary": "str ā€” overall performance analysis", + "estimated_improvement": "str ā€” expected speedup (e.g. '10x faster')", + "approved": "bool ā€” True if already optimal", + }, + } + for t in TASKS.values() + ] diff --git a/tests/envs/test_sql_optim_environment.py b/tests/envs/test_sql_optim_environment.py new file mode 100644 index 000000000..1fe1bbe1f --- /dev/null +++ b/tests/envs/test_sql_optim_environment.py @@ -0,0 +1,384 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Tests for the SQL Query Optimization environment. + +Covers the server-side environment (reset/step/state driven against real +DuckDB execution), the wire models, and the client's payload/parse helpers. +The client's transport is exercised separately by the WebSocket integration +tests; here we unit-test the pure parsing methods without a live server. +""" + +import os +import sys + +import pytest + +# Add the repo root so `envs.*` imports resolve. +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) + +from envs.sql_optim_env import ( # noqa: E402 + SQLOptimAction, + SQLOptimEnv, + SQLOptimObservation, + SQLOptimState, +) +from envs.sql_optim_env.server.sql_optim_environment import ( # noqa: E402 + DEFAULT_TASK_ID, + SQLOptimEnvironment, +) +from envs.sql_optim_env.server.tasks import TASKS # noqa: E402 + + +@pytest.fixture(scope="module") +def env(): + """A single environment instance (DuckDB tables are built once).""" + return SQLOptimEnvironment() + + +@pytest.fixture(scope="module") +def executor(): + """The shared DuckDB query executor (tables built once).""" + from envs.sql_optim_env.server.executor import get_executor + + return get_executor() + + +def _good_action() -> SQLOptimAction: + """A result-preserving task-1 rewrite: sargable predicates, projection kept.""" + return SQLOptimAction( + suggestions=[ + { + "issue_type": "select_star", + "line": 1, + "description": "SELECT * is wasteful", + "severity": "high", + "fix": "project columns", + }, + { + "issue_type": "non_sargable_cast", + "line": 3, + "description": "CAST blocks pruning", + "severity": "high", + "fix": "compare int", + }, + { + "issue_type": "function_on_date_column", + "line": 4, + "description": "year() non-sargable", + "severity": "medium", + "fix": "date range", + }, + ], + optimized_query=( + "SELECT * FROM orders WHERE customer_id = 5000 " + "AND created_at >= DATE '2024-01-01' AND created_at < DATE '2025-01-01';" + ), + summary="Made the CAST and year() predicates sargable; flagged SELECT *.", + estimated_improvement="~10x faster", + approved=False, + ) + + +class TestReset: + def test_reset_returns_task_observation(self, env): + obs = env.reset(task_id="task_1_basic_antipatterns") + assert isinstance(obs, SQLOptimObservation) + assert obs.task_id == "task_1_basic_antipatterns" + assert obs.step_count == 0 + assert obs.reward is None + assert obs.done is False + assert obs.sql_query # a query is presented + + def test_reset_defaults_to_first_task(self, env): + obs = env.reset() + assert obs.task_id == DEFAULT_TASK_ID + + def test_reset_unknown_task_raises(self, env): + with pytest.raises(ValueError, match="Unknown task_id"): + env.reset(task_id="does_not_exist") + + +class TestStep: + def test_step_rewards_a_faster_correct_rewrite(self, env): + env.reset(task_id="task_1_basic_antipatterns") + obs = env.step(_good_action()) + assert isinstance(obs.reward, float) + assert obs.reward > 0.0 + assert "reward_breakdown" in obs.metadata + assert obs.metadata["execution"]["speedup"] > 1.0 # real DuckDB speedup + + def test_step_accumulates_issue_context(self, env): + env.reset(task_id="task_1_basic_antipatterns") + obs = env.step(_good_action()) + assert set(obs.issues_found_so_far) >= {"select_star", "non_sargable_cast"} + + def test_episode_terminates_at_max_steps(self, env): + env.reset(task_id="task_1_basic_antipatterns") + max_steps = TASKS["task_1_basic_antipatterns"]["max_steps"] + empty = SQLOptimAction(optimized_query="") # no reward, forces max_steps + done = False + for _ in range(max_steps): + done = env.step(empty).done + assert done is True + + def test_step_after_done_raises(self, env): + env.reset(task_id="task_1_basic_antipatterns") + max_steps = TASKS["task_1_basic_antipatterns"]["max_steps"] + empty = SQLOptimAction(optimized_query="") + for _ in range(max_steps): + env.step(empty) + with pytest.raises(RuntimeError, match="Episode finished"): + env.step(empty) + + +class TestState: + def test_state_tracks_episode(self, env): + env.reset(task_id="task_2_correlated_subqueries") + state = env.state + assert isinstance(state, SQLOptimState) + assert state.step_count == 0 + assert state.episode_done is False + assert state.cumulative_reward == 0.0 + + +class TestModels: + def test_action_roundtrips(self): + action = _good_action() + dumped = action.model_dump() + assert dumped["optimized_query"].startswith("SELECT *") + assert SQLOptimAction(**dumped) == action + + def test_observation_carries_base_fields(self): + obs = SQLOptimObservation(task_id="t", reward=0.5, done=True) + assert obs.reward == 0.5 and obs.done is True + + +class TestClientParsing: + """The client's pure helpers, without opening a WebSocket.""" + + def _client(self): + return SQLOptimEnv(base_url="http://localhost:8000") + + def test_step_payload(self): + payload = self._client()._step_payload(_good_action()) + assert payload["optimized_query"].startswith("SELECT *") + assert payload["approved"] is False + + def test_parse_result(self): + client = self._client() + result = client._parse_result( + { + "observation": {"task_id": "t", "sql_query": "SELECT 1"}, + "reward": 0.42, + "done": True, + "metadata": {"feedback": "ok"}, + } + ) + assert isinstance(result.observation, SQLOptimObservation) + assert result.reward == 0.42 + assert result.done is True + assert result.metadata == {"feedback": "ok"} + + def test_parse_state(self): + state = self._client()._parse_state({"task_id": "t", "step_count": 3}) + assert isinstance(state, SQLOptimState) + assert state.task_id == "t" and state.step_count == 3 + + +class TestExecutionSafetyAndCorrectness: + """Regression tests for the execution-engine review findings. + + Cover DB isolation, checksum robustness, single-compare-per-step, stale + execution on empty rewrites, and the task-definition corrections. + """ + + def test_mutating_rewrite_is_rejected_and_db_intact(self, executor): + """A DROP/DELETE rewrite is rejected and never mutates the shared DB.""" + before = executor.table_stats["orders"] + for bad in ( + "DROP TABLE orders", + "DELETE FROM orders", + "SELECT 1; DROP TABLE orders", + ): + res = executor.compare("SELECT 1", bad) + assert res["optimized_error"], f"not rejected: {bad}" + assert res["results_match"] is False + assert executor.table_stats["orders"] == before # nothing persisted + + def test_read_only_rewrite_still_runs(self, executor): + """A legitimate SELECT rewrite executes normally.""" + res = executor.compare( + "SELECT COUNT(*) FROM orders", "SELECT COUNT(*) FROM orders" + ) + assert res["optimized_error"] is None + assert res["results_match"] is True + + def test_trailing_semicolon_large_result_checksum(self, executor): + """A `;`-terminated query over >50k rows uses the checksum path correctly.""" + q = "SELECT * FROM events;" # 1M rows -> checksum branch, trailing ';' + res = executor.compare(q, q) + assert res["results_match"] is True + assert res["original_rows"] == 1_000_000 + + def test_task3_equality_rewrite_is_result_preserving(self, executor): + """Task 3's intended equality rewrite returns exactly the same rows. + + (Speedup is not asserted ā€” wall-clock timing is too noisy for a stable + threshold; result-preservation is the property that matters here.) + """ + original = TASKS["task_3_wildcard_scan"]["sql_query"] + rewrite = original.replace( + "WHERE event_type LIKE '%purchase%'\n OR event_type LIKE '%buy%';", + "WHERE event_type = 'purchase';", + ) + assert rewrite != original # the replacement actually happened + res = executor.compare(original, rewrite) + assert res["results_match"] is True + assert res["optimized_error"] is None + + def test_task5_is_deterministic(self, executor): + """Task 5's windows are deterministic, so a result-preserving rewrite can match.""" + q = TASKS["task_5_window_functions"]["sql_query"] + assert executor.compare(q, q)["results_match"] is True + + def test_task4_subqueries_labeled_uncorrelated(self): + """Task 4 no longer mislabels its uncorrelated scalar subqueries as correlated.""" + types = [ + g["type"] for g in TASKS["task_4_implicit_join"]["ground_truth_issues"] + ] + assert "uncorrelated_scalar_subquery_avg" in types + assert "uncorrelated_scalar_subquery_max" in types + blob = str(TASKS["task_4_implicit_join"]).lower() + assert "uncorrelated" in blob + assert "per group" not in blob # the "recomputed per group" mislabel is gone + + def test_step_runs_compare_once(self, monkeypatch): + """`step` reuses the grader's comparison instead of running DuckDB twice.""" + from envs.sql_optim_env.server import executor as ex_mod + + ex = ex_mod.get_executor() + real = ex.compare + calls = {"n": 0} + + def counting(*a, **k): + calls["n"] += 1 + return real(*a, **k) + + monkeypatch.setattr(ex, "compare", counting) + e = SQLOptimEnvironment() + e.reset(task_id="task_1_basic_antipatterns") + e.step(_good_action()) + assert calls["n"] == 1 # was 2 before the dedup fix + + def test_empty_rewrite_does_not_report_stale_execution(self): + """An empty rewrite clears `execution` instead of keeping the prior step's.""" + e = SQLOptimEnvironment() + e.reset(task_id="task_1_basic_antipatterns") + obs1 = e.step( + SQLOptimAction( + suggestions=[{"issue_type": "select_star", "severity": "low"}], + optimized_query="SELECT id FROM orders WHERE customer_id = 5000", + summary="partial", + approved=False, + ) + ) + assert obs1.metadata["execution"] is not None + assert obs1.done is False + obs2 = e.step( + SQLOptimAction(optimized_query="", summary="no rewrite", approved=False) + ) + assert obs2.metadata["execution"] is None # not the stale step-1 comparison + + def test_valid_rewrites_run_including_comments_and_literals(self, executor): + """Valid read-only rewrites execute ā€” no structural pre-check false-rejects. + + The read-only connection is the only gate, so rewrites using function + names (`REPLACE`), string literals containing `set`/`;`, or a leading SQL + comment all run instead of being rejected up front. + """ + for good in ( + "SELECT REPLACE(status, 'a', 'b') FROM orders LIMIT 1", + "SELECT COUNT(*) FROM orders WHERE status = 'set'", + "SELECT COUNT(*) FROM orders WHERE status = 'a;b'", + "-- optimized\nSELECT COUNT(*) FROM orders", + ): + res = executor.compare("SELECT COUNT(*) FROM orders", good) + assert res["optimized_error"] is None, f"falsely rejected: {good!r}" + + def test_issue_detection_ignores_the_rewrite_sql(self): + """Echoing the slow SQL earns no issue-detection credit without analysis. + + Anti-pattern tokens live in the query itself; scoring detection off the + rewrite would let an agent farm the slice by copying the original query. + """ + from envs.sql_optim_env.server.graders import grade + + original = TASKS["task_1_basic_antipatterns"]["sql_query"] + # Rewrite echoes the original SQL, but the analysis fields are empty. + action = SQLOptimAction( + suggestions=[], optimized_query=original, summary="", approved=False + ) + reward = grade(TASKS["task_1_basic_antipatterns"], action) + assert reward.breakdown["issue_detection"] == 0.0 + + def test_read_only_connection_rejects_writes_at_engine_level(self, executor): + """Even a DML-in-CTE form is rejected ā€” the connection is truly read-only.""" + before = executor.table_stats["orders"] + res = executor.compare("SELECT 1", "WITH d AS (SELECT 1) DELETE FROM orders") + assert res["optimized_error"] # engine (or structural gate) rejects it + assert executor.table_stats["orders"] == before + + def test_task1_result_preserving_rewrite_matches(self, executor): + """Task 1's intended rewrite (sargable predicates, SELECT * kept) matches.""" + original = TASKS["task_1_basic_antipatterns"]["sql_query"] + rewrite = ( + "SELECT * FROM orders WHERE customer_id = 5000 " + "AND created_at >= DATE '2024-01-01' AND created_at < DATE '2025-01-01'" + ) + res = executor.compare(original, rewrite) + assert res["results_match"] is True + + def test_task2_last_order_is_deterministic(self, executor): + """Task 2's `last_order_amount` subquery has a unique tie-breaker.""" + q = TASKS["task_2_correlated_subqueries"]["sql_query"] + assert executor.compare(q, q)["results_match"] is True + + def test_speedup_requires_correctness(self): + """A fast but wrong rewrite earns no speedup credit (only when results match).""" + from envs.sql_optim_env.server.graders import grade + + task = TASKS["task_5_window_functions"] + # Trivially fast, but wrong results. + wrong = SQLOptimAction(optimized_query="SELECT 1", summary="x" * 60) + r_wrong = grade(task, wrong) + assert r_wrong.breakdown["execution_speedup"] == 0.0 + assert r_wrong.execution["results_match"] is False + + def test_concurrent_compare_is_safe(self, executor): + """Concurrent access to the shared connection does not corrupt or crash.""" + import threading + + errors = [] + + def worker(): + try: + for _ in range(3): + executor.compare( + "SELECT COUNT(*) FROM orders", + "SELECT COUNT(*) FROM orders WHERE status = 'completed'", + ) + assert executor.table_stats["orders"] == 500_000 + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(6)] + for t in threads: + t.start() + for t in threads: + t.join() + assert not errors From 8f0133b6dafb47d1df7abfedd26c34874312d94c Mon Sep 17 00:00:00 2001 From: OfficialAbhinavSingh Date: Sat, 25 Jul 2026 23:08:39 +0530 Subject: [PATCH 2/4] fix(sql-optim-env): measure queries inside DuckDB instead of materializing rows _run() fetched every result row into Python via fetchall(), three times over for the timing median. For a query like SELECT * FROM events (1M rows) that held the whole result in memory and made the measurement report tuple construction rather than query execution: fetchall() took 783.6ms of which DuckDB execution was 25.5ms, so 97% of the reported time was client-side conversion. Timing now comes from DuckDB's profiler (EXPLAIN ANALYZE -> Total Time), which executes the query for real and discards the result inside the engine, so no rows cross into Python. Rows are materialized only up to the 50k cap compare() already used to pick its precise row-by-row comparison; above that the exact row count comes from an in-engine COUNT(*) and correctness from the existing order-independent checksum. A wall-clock streaming drain is kept as a fallback for a DuckDB version whose profile footer cannot be parsed, so timing degrades instead of breaking. The execution lock becomes reentrant so a whole measurement stays serialized as before. Comparing SELECT * FROM events goes from a 867.7MB peak and 4212ms reported to 14.5MB and 18.5ms; task_5_window_functions compare() drops from 13.6s to 5.7s. Small-result queries pay one extra profiled execution (task_1 101ms -> 233ms); the total across all five tasks is still ~45% lower. Behaviour is otherwise unchanged: writes are still rejected by the read-only connection, row counts stay exact, and results under the cap are still compared row by row. --- envs/sql_optim_env/server/executor.py | 195 +++++++++++++++++++---- tests/envs/test_sql_optim_environment.py | 94 +++++++++++ 2 files changed, 261 insertions(+), 28 deletions(-) diff --git a/envs/sql_optim_env/server/executor.py b/envs/sql_optim_env/server/executor.py index bc543dcc6..c75acde8a 100644 --- a/envs/sql_optim_env/server/executor.py +++ b/envs/sql_optim_env/server/executor.py @@ -14,6 +14,7 @@ import atexit import os +import re import shutil import tempfile import threading @@ -25,6 +26,21 @@ _instance: Optional["QueryExecutor"] = None _lock = threading.Lock() +# Result sets larger than this are never pulled into Python: the row count and +# the correctness check are computed inside DuckDB instead. `SELECT * FROM +# events` is a million rows, and materializing that (three times, for the +# timing median) both dominated the measurement and held the whole result in +# memory. The cap matches the threshold `compare()` already used to pick the +# precise row-by-row comparison, so behaviour for small results is unchanged. +_MAX_MATERIALIZED_ROWS = 50_000 +_FETCH_BATCH = 10_000 +_TIMING_RUNS = 3 +# `EXPLAIN ANALYZE` prints seconds with four decimals, so 0.1 ms is the smallest +# value it can express. Used as a floor so a sub-resolution query can never make +# the speedup ratio divide by zero. +_MIN_MEASURABLE_MS = 0.1 +_TOTAL_TIME_RE = re.compile(r"Total Time:\s*([0-9]*\.?[0-9]+)\s*s") + def _strip_terminators(query: str) -> str: """Trim surrounding whitespace and any trailing semicolons. @@ -57,7 +73,9 @@ def __init__(self) -> None: builder.close() self.conn = duckdb.connect(self._path, read_only=True, config={"threads": "2"}) - self._exec_lock = threading.Lock() + # Reentrant so `_run` can hold the lock across a whole measurement while + # the helpers it calls still take it individually. + self._exec_lock = threading.RLock() atexit.register(self._cleanup) def _cleanup(self) -> None: @@ -138,28 +156,145 @@ def _build_tables(self, conn: "duckdb.DuckDBPyConnection") -> None: # ā”€ā”€ Execution helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ - def _run( - self, query: str, runs: int = 3 - ) -> Tuple[float, Optional[List], Optional[str]]: + def _probe(self, query: str) -> Tuple[Optional[List], Optional[str]]: + """ + Fetch at most ``_MAX_MATERIALIZED_ROWS`` rows of *query*. + + Returns (rows, error). ``rows`` is None when the result set is larger + than the cap ā€” the caller then works from the in-engine row count and + checksum instead of holding the full result in Python. + + This is also the authoritative error check: the query runs unwrapped on + the read-only connection, so a write (DDL/DML, inside a CTE, or after a + ``;``) still surfaces as a DuckDB error exactly as before. + """ + try: + with self._exec_lock: + # `execute()` returns the connection itself, so the leftover rows + # of an over-cap result cannot be released with a `close()` + # (that would close the shared read-only connection). The next + # statement on the connection supersedes the pending result, + # and every caller issues one straight after. + head = self.conn.execute(query).fetchmany(_MAX_MATERIALIZED_ROWS + 1) + except Exception as exc: + return None, str(exc) + if len(head) > _MAX_MATERIALIZED_ROWS: + return None, None + return head, None + + def _row_count(self, query: str) -> Tuple[Optional[int], Optional[str]]: + """ + Exact row count for *query*, computed inside DuckDB ā€” no rows cross into + Python. Falls back to a streaming drain (bounded to one batch of memory) + for the rare query that cannot be wrapped in a subquery. """ - Execute *query* up to *runs* times on the read-only connection. - Returns (median_ms, rows, error_or_None). A write raises and is returned - as the error string. + try: + with self._exec_lock: + return ( + self.conn.execute(f"SELECT COUNT(*) FROM ({query}) t").fetchone()[ + 0 + ], + None, + ) + except Exception: + pass + + try: + count = 0 + with self._exec_lock: + cursor = self.conn.execute(query) + while True: + batch = cursor.fetchmany(_FETCH_BATCH) + if not batch: + break + count += len(batch) + return count, None + except Exception as exc: + return None, str(exc) + + def _engine_ms(self, query: str, runs: int) -> Optional[float]: + """ + Median DuckDB-side execution time in ms, or None if unavailable. + + ``EXPLAIN ANALYZE`` really executes the query and discards the result + inside the engine, so this measures execution alone: no rows are + converted to Python objects and nothing is buffered client-side. Returns + None when the profile footer cannot be parsed (a DuckDB version whose + output differs), which makes the caller fall back to wall-clock timing. """ timings: List[float] = [] - rows: Optional[List] = None + for _ in range(runs): + try: + with self._exec_lock: + rows = self.conn.execute(f"EXPLAIN ANALYZE {query}").fetchall() + except Exception: + return None + match = _TOTAL_TIME_RE.search("\n".join(str(r[-1]) for r in rows)) + if match is None: + return None + timings.append(float(match.group(1)) * 1000.0) - with self._exec_lock: - for _ in range(runs): - try: + timings.sort() + return round(max(timings[len(timings) // 2], _MIN_MEASURABLE_MS), 3) + + def _wall_ms(self, query: str, runs: int) -> Tuple[float, Optional[str]]: + """ + Wall-clock fallback timing: time a streaming drain that discards rows. + + Includes client-side conversion cost, so it is less faithful than + ``_engine_ms``, but memory stays bounded to a single batch instead of + the whole result set. + """ + timings: List[float] = [] + for _ in range(runs): + try: + with self._exec_lock: t0 = time.perf_counter() - rows = self.conn.execute(query).fetchall() + cursor = self.conn.execute(query) + while cursor.fetchmany(_FETCH_BATCH): + pass timings.append((time.perf_counter() - t0) * 1000.0) - except Exception as exc: - return 99_999.0, None, str(exc) + except Exception as exc: + return 99_999.0, str(exc) timings.sort() - return round(timings[len(timings) // 2], 3), rows, None + return round(max(timings[len(timings) // 2], _MIN_MEASURABLE_MS), 3), None + + def _run( + self, query: str, runs: int = _TIMING_RUNS + ) -> Tuple[float, Optional[int], Optional[List], Optional[str]]: + """ + Execute *query* on the read-only connection and measure it. + + Returns (median_ms, row_count, rows, error_or_None). ``row_count`` is + exact whenever there is no error; ``rows`` is None for result sets above + ``_MAX_MATERIALIZED_ROWS``. A write raises and is returned as the error + string. + + The timing comes from DuckDB's own profiler rather than from wrapping a + ``fetchall()``, so it reflects query execution instead of the cost of + building Python tuples ā€” for a 1M-row scan the latter was roughly 30x + the former and swamped the signal the reward function is built on. + """ + with self._exec_lock: + rows, error = self._probe(query) + if error is not None: + return 99_999.0, None, None, error + + if rows is not None: + row_count: Optional[int] = len(rows) + else: + row_count, error = self._row_count(query) + if error is not None: + return 99_999.0, None, None, error + + median_ms = self._engine_ms(query, runs) + if median_ms is None: + median_ms, error = self._wall_ms(query, runs) + if error is not None: + return 99_999.0, None, None, error + + return median_ms, row_count, rows, None def _checksum( self, query: str @@ -219,28 +354,31 @@ def compare(self, original: str, optimized: str) -> Dict[str, Any]: original = _strip_terminators(original) optimized = _strip_terminators(optimized) - orig_ms, orig_rows, orig_err = self._run(original) + orig_ms, orig_count, orig_rows, orig_err = self._run(original) # The optimized query is agent-authored, but the read-only connection is # the guarantee: any write (DDL/DML, in a CTE, or after a `;`) is rejected # by the DuckDB engine and surfaces as `optimized_error`. No structural # pre-check is needed, which also avoids false rejections of valid # rewrites (leading comments, a `;` inside a string literal, etc.). - opt_ms, opt_rows, opt_err = self._run(optimized) + opt_ms, opt_count, opt_rows, opt_err = self._run(optimized) # ā”€ā”€ Correctness: do both queries return the same data? ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ # Use a DuckDB-level checksum (order-independent) to avoid # false negatives from non-deterministic row ordering in parallel # window function queries on large tables. results_match = False - if orig_rows is not None and opt_rows is not None: + if orig_err is None and opt_err is None: try: - if len(orig_rows) != len(opt_rows): + if orig_count != opt_count: results_match = False - elif len(orig_rows) == 0: + elif orig_count == 0: results_match = True - elif len(orig_rows) <= 50_000: - # Small/medium: full sorted comparison (precise) + elif orig_rows is not None and opt_rows is not None: + # Small/medium: full sorted comparison (precise). Equal row + # counts mean both sides are on the same side of the + # materialization cap, so either both rows are present or + # neither is. orig_s = sorted(str(r) for r in orig_rows) opt_s = sorted(str(r) for r in opt_rows) results_match = orig_s == opt_s @@ -249,13 +387,14 @@ def compare(self, original: str, optimized: str) -> Dict[str, Any]: # (deterministic regardless of row ordering / thread count) o_cnt, o_chk, o_err2 = self._checksum(original) p_cnt, p_chk, p_err2 = self._checksum(optimized) - if o_err2 or p_err2: - # Checksum failed ā€” fall back to row count - results_match = len(orig_rows) == len(opt_rows) + if o_err2 or p_err2 or o_chk is None or p_chk is None: + # No usable checksum: the equal row counts established + # above are the only signal left, as before. + results_match = True else: results_match = (o_cnt == p_cnt) and (o_chk == p_chk) except Exception: - results_match = len(orig_rows) == len(opt_rows) + results_match = orig_count == opt_count # ā”€ā”€ Speedup ratio ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ speedup = 1.0 @@ -283,8 +422,8 @@ def compare(self, original: str, optimized: str) -> Dict[str, Any]: "optimized_ms": opt_ms, "speedup": speedup, "results_match": results_match, - "original_rows": len(orig_rows) if orig_rows is not None else 0, - "optimized_rows": len(opt_rows) if opt_rows is not None else 0, + "original_rows": orig_count if orig_count is not None else 0, + "optimized_rows": opt_count if opt_count is not None else 0, "original_error": orig_err, "optimized_error": opt_err, "verdict": verdict, diff --git a/tests/envs/test_sql_optim_environment.py b/tests/envs/test_sql_optim_environment.py index 1fe1bbe1f..4148c5ed2 100644 --- a/tests/envs/test_sql_optim_environment.py +++ b/tests/envs/test_sql_optim_environment.py @@ -15,6 +15,8 @@ import os import sys +import time +import tracemalloc import pytest @@ -382,3 +384,95 @@ def worker(): for t in threads: t.join() assert not errors + + +class TestLargeResultMeasurement: + """Large result sets are measured without being pulled into Python. + + Timing comes from DuckDB's profiler and correctness from the in-engine + checksum, so a million-row query neither buys a million Python tuples nor + reports the cost of building them as its execution time. + """ + + def test_large_result_is_not_materialized(self, executor): + """A 1M-row query yields an exact count and no materialized rows.""" + ms, count, rows, err = executor._run("SELECT * FROM events") + assert err is None + assert rows is None # never held in Python + assert count == 1_000_000 # still exact, computed inside DuckDB + assert ms < 90_000 # a real measurement, not the error sentinel + + def test_small_result_is_still_compared_row_by_row(self, executor): + """Below the cap, rows are materialized and compared precisely.""" + ms, count, rows, err = executor._run("SELECT * FROM products") + assert err is None + assert count == 1_000 + assert rows is not None and len(rows) == 1_000 + + # Precise comparison, not count-only: same rows in a different order + # match, while an equal-sized result with different values does not. + reordered = executor.compare( + "SELECT id, price FROM products ORDER BY id", + "SELECT id, price FROM products ORDER BY id DESC", + ) + assert reordered["results_match"] is True + altered = executor.compare( + "SELECT id, price FROM products", + "SELECT id, price + 1 AS price FROM products", + ) + assert altered["results_match"] is False + assert altered["original_rows"] == altered["optimized_rows"] == 1_000 + + def test_timing_excludes_python_materialization(self, executor): + """The reported time tracks DuckDB execution, not tuple construction.""" + query = "SELECT * FROM events" + start = time.perf_counter() + with executor._exec_lock: + fetched = executor.conn.execute(query).fetchall() + fetchall_ms = (time.perf_counter() - start) * 1000.0 + assert len(fetched) == 1_000_000 + del fetched + + reported_ms = executor._run(query)[0] + # Scanning the table costs a few percent of what turning it into a + # million Python tuples does, so this bound has an order of magnitude of + # headroom for a slow CI machine while still failing outright if + # `fetchall()` timing comes back. + assert reported_ms < fetchall_ms * 0.25 + + def test_peak_memory_is_bounded_for_large_results(self, executor): + """Comparing a 500k-row query does not buffer the whole result set.""" + tracemalloc.start() + try: + result = executor.compare("SELECT * FROM orders", "SELECT * FROM orders") + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + assert result["results_match"] is True + assert result["original_rows"] == 500_000 + # Materializing 500k rows twice cost ~430 MB before; the cap keeps this + # to a couple of batches regardless of result size. + assert peak < 100 * 1024 * 1024 + + def test_timing_is_floored_at_the_profiler_resolution(self, executor): + """A query faster than the profiler can express still gets a usable time. + + `EXPLAIN ANALYZE` prints seconds to four decimals, so it reports 0.0000s + for a query answered from table metadata. Guards the invariant rather + than reproducing a past failure: every reported time stays at or above + the floor, so the speedup ratio can never divide by zero. + """ + from envs.sql_optim_env.server.executor import _MIN_MEASURABLE_MS + + for query in ( + "SELECT COUNT(*) FROM events", # answered from metadata + "SELECT 1", + "SELECT * FROM products WHERE id = -1", # empty result + ): + reported_ms = executor._run(query)[0] + assert reported_ms >= _MIN_MEASURABLE_MS, query + result = executor.compare(query, query) + assert result["optimized_ms"] >= _MIN_MEASURABLE_MS + assert result["speedup"] > 0 + assert result["results_match"] is True From 2bf81ee8b77faf96486c6c2148e5d0e25bf9aa52 Mon Sep 17 00:00:00 2001 From: OfficialAbhinavSingh Date: Sat, 25 Jul 2026 23:28:05 +0530 Subject: [PATCH 3/4] fix(sql-optim-env): sandbox agent SQL from the filesystem and stop unverified match credit Two High-severity review findings on the execution engine. The read-only connection stops the agent rewrite from changing the database, but not from reaching the filesystem. On the previous commit an agent-authored query could run `COPY ... TO` to write an arbitrary file (verified: the file was created), `read_csv('/etc/passwd')` to read one, plus `read_text`, `read_blob`, `read_parquet`, `glob`, `COPY FROM`, `EXPORT DATABASE`, `ATTACH` and `INSTALL`/`LOAD`. Both connections now open with `enable_external_access=false`, and with `lock_configuration=true` so a rewrite cannot `SET` or `PRAGMA` its way back out (DuckDB additionally refuses to re-enable external access on a running database). The escape hatches are blocked on the `EXPLAIN ANALYZE` timing path as well. Second, a rewrite ending in a `--` comment silently bought full correctness credit. The checksum wraps the query in a subquery, and on a single line the trailing comment swallowed the closing parenthesis, so the checksum failed and equal row counts alone were reported as a match: `SELECT id FROM events -- c` against `SELECT id + 1 AS id FROM events -- c` matched on a million rows that all differed, unlocking the speedup slice too. `_as_subquery()` now puts the query on its own line, and an over-cap result with no usable checksum reports `results_match=False` instead of guessing, as does the surrounding except path. Note that the review's stated mechanism for the second finding -- `str.format` raising on braces in the query -- does not reproduce: `format()` does not rescan substituted values, and a brace-containing query checksums normally. The trailing-comment path is the one that was exploitable. TestSandboxAndCorrectnessCredit covers both; five of its six tests fail on the previous commit. 38 env tests pass, 1544 repo-wide, ruff clean. --- envs/sql_optim_env/server/executor.py | 70 +++++++++++------ tests/envs/test_sql_optim_environment.py | 96 ++++++++++++++++++++++++ 2 files changed, 141 insertions(+), 25 deletions(-) diff --git a/envs/sql_optim_env/server/executor.py b/envs/sql_optim_env/server/executor.py index c75acde8a..d9b9ccb12 100644 --- a/envs/sql_optim_env/server/executor.py +++ b/envs/sql_optim_env/server/executor.py @@ -41,16 +41,39 @@ _MIN_MEASURABLE_MS = 0.1 _TOTAL_TIME_RE = re.compile(r"Total Time:\s*([0-9]*\.?[0-9]+)\s*s") +# Opening the connection read-only stops the agent from changing the *database*, +# but on its own it leaves DuckDB's filesystem escape hatches open: `COPY ... TO` +# writes arbitrary files, `read_csv`/`read_text`/`glob` read them, `ATTACH` +# mounts other databases and `INSTALL`/`LOAD` pull in extensions. Since the query +# under measurement is agent-authored, external access is disabled outright, and +# the configuration is locked so a rewrite cannot `SET` its way back out (DuckDB +# also refuses to re-enable external access on a running database). +_DUCKDB_CONFIG = { + "threads": "2", + "enable_external_access": "false", + "lock_configuration": "true", +} + def _strip_terminators(query: str) -> str: """Trim surrounding whitespace and any trailing semicolons. Trailing ``;`` breaks queries that are wrapped in a subquery - (``SELECT COUNT(*) FROM ({query}) t``) for correctness checksums. + (``SELECT COUNT(*) FROM () t``) for correctness checksums. """ return query.strip().rstrip(";").strip() +def _as_subquery(query: str, select: str) -> str: + """Build ``SELECT