Project 03 · AI-assisted data engineering
AI-Assisted SQL Optimizer
A CLI that uses Claude to suggest Spark and Snowflake query rewrites and partition strategies. Benchmarked against a seeded 5-category corpus with sqlglot-driven heuristics and an optional EXPLAIN verifier.
Measured 2026-07 on the local demo path (synthetic data, seed 42) — reproduce via the repo's Results section.
Live demo
Pick a query, see the rewrite
Output is replayed from a static fixture so the page never makes an external API call by default. The repo includes the full CLI plus an optional Cloudflare Worker proxy for self-hosted live mode.
Live demo
Optimize a query
Pick one of the five corpus samples. The original goes in; Claude returns a rewrite, reasoning, and an estimated cost delta. Output is replayed from a fixture so the page never makes an API call by default.
Original SQL
Inline diff
original vs optimizedWhy this rewrite
Claude’s reasoning
- Push country and date filters into subqueries so each side scans fewer rows before the join.
- On Snowflake / Iceberg this lets the engine prune partitions and apply Bloom filters before shuffle.
- Typical 3-10x reduction in shuffle bytes on joins with selective filters.
Pattern note · Filters higher in the plan get pushed down to scan-time, dropping rows before they participate in the join.
Corpus
Estimated cost reduction across the sample queries
Per-query cost-reduction estimates from the seeded benchmark corpus. Each of the five corpus queries is shown here with its ground truth in the repo.
Estimated cost reduction by query
Illustrative cost-reduction estimates — see methodology below
Numbers reflect planner heuristics on the seeded benchmark corpus, not production runs. EXPLAIN-based scoring noted in the repo.
| Query | Estimated cost reduction |
|---|---|
| Predicate pushdown into the inner join | 62% |
| Broadcast small dimension instead of shuffle join | 45% |
| Add partition predicate so the planner can prune | 88% |
| Replace self-join with window function | 70% |
| Project only required columns | 30% |
Methodology
How the optimizer works, scores, and where it fails
Three columns: how it works, how suggestions are scored, and the limitations a senior reviewer would call out.
How it works
- sqlglot parses the input into an AST and a small heuristic pass collects findings (function-wrapped partition columns, SELECT *, exact COUNT DISTINCT, missing broadcast hints).
- Findings + raw SQL are sent to the Anthropic Messages API with a structured optimizer prompt. Output is constrained to a JSON envelope of rewrite + reasoning.
- An optional EXPLAIN runner (Spark or Snowflake) executes both plans against a shadow database and reports the cost delta back to the CLI.
Scoring rubric
- Keyword overlap with ground-truth labels in
corpus/ground_truth.yaml— coarse but cheap to run on every change. Measured: 0.67 average in dry-run mode. - EXPLAIN cost delta when a Spark or Snowflake engine is reachable; treated as approximate, not authoritative.
- Human review across the corpus as it grows, recorded in the repo with the reviewer’s rationale per query.
Limitations
- Claude can hallucinate column or table names not in the schema — especially for novel dialects.
- EXPLAIN-based scoring is approximate: a "cheaper" plan can still be slower in practice.
- Dialect drift between Spark SQL and Snowflake catches edge cases (QUALIFY, lateral views, named struct fields).
- Suggestions require human review before applying to production. The CLI defaults to printing diffs, never auto-rewriting files.
Local vs. live
Runs locally by default. Claude is optional.
The on-page demo replays pre-computed responses, so no API keys
are bundled into the client and the default build never reaches
the network. The repo ships the full CLI with a heuristic /
mock-response mode that works offline, plus an optional
Cloudflare Worker proxy at
src/workers/sql-optimizer-proxy/
for live calls when self-hosted with your own
ANTHROPIC_API_KEY. When
you deploy the Worker, set
PUBLIC_LIVE_DEMO_URL at
build time and the demo swaps fixture lookup for a Worker fetch.
Architecture
How it works under the hood
Sequence diagram of the full flow, including the optional EXPLAIN-runner branch. The engineering narrative below covers the judgment calls.
sequenceDiagram
participant U as User
participant CLI as sql-optimizer CLI
participant A as Analyzer (sqlglot + heuristics)
participant C as Claude API
participant E as EXPLAIN runner (optional)
U->>CLI: analyze query.sql
CLI->>A: parse + collect findings
A-->>CLI: AST + findings
CLI->>C: optimizer_prompt + sql + findings
C-->>CLI: rewrite + reasoning
CLI->>E: EXPLAIN original / EXPLAIN rewrite
E-->>CLI: cost deltas
CLI-->>U: markdown diff + cost report
Engineering narrative
Problem, judgment calls, and measured results
The write-up a senior reviewer would ask for — why this architecture, what was traded away, and what the numbers actually say.
Problem
There is a real gap between “junior writes a working query” and “senior knows when to broadcast, partition-prune, flatten a CTE, or rewrite a window function.” This tool closes part of that gap by sending parsed SQL plus heuristic findings to Claude with a structured optimizer prompt, then optionally verifying the rewrite against a real EXPLAIN plan. It targets Spark SQL and Snowflake, the two dialects I work with most.
Judgment calls
Heuristics + LLM, not LLM only. Sending raw SQL to Claude works, but
the model spends most of its budget re-deriving facts a parser already
knows: which columns appear in the projection, where the partition column
lives, whether the predicate is wrapped in a function. A small sqlglot pass
extracts those cheaply and ships them as structured findings, leaving the
model to decide what to do with them - and keeping the offline --dry-run
mode honest.
Why EXPLAIN is optional. EXPLAIN cost numbers are informative, not ground truth - they reflect the planner’s stale statistics and its own broadcast threshold. The CLI treats them as a tiebreaker between candidate rewrites, never as the deciding signal.
Prompt sensitivity is a real risk. Small changes to the optimizer prompt materially change suggestion quality. The prompt template is versioned in the repo, and corpus tests pin a known prompt version so regressions surface in CI rather than in production.
A corpus designed to grow. Five optimization categories, one fully
specified query per category, each with explicit ground truth in
corpus/ground_truth.yaml. Adding a query is one .sql file plus one
ground-truth entry - the benchmark discovers both automatically. A small
corpus that matches its documentation beats a padded one.
Measured results
make benchmark in heuristic-only dry-run mode scores 0.67 average keyword
overlap against ground truth, with correct findings on 4 of 5 queries at
~68 ms analyzer latency per query. The one miss is instructive: the
cte_flattening query’s rewrite vocabulary (“single scan”, “CASE”) only
appears once Claude proposes the rewrite - the heuristics flag the multi-scan
pattern but do not name the fix. That gap is exactly what the LLM adds.
Human review and EXPLAIN cost deltas are specified in the methodology but not
yet run - the README keeps them under “Planned evaluation” rather than
pretending.
Honest limitations
- Claude can hallucinate column or table names that look plausible but do not exist in the schema.
- EXPLAIN-based scoring is only as good as the optimizer’s cost estimates - a “cheaper” plan can still be slower in practice.
- Dialect drift between Spark SQL and Snowflake catches edge cases (e.g.
QUALIFY, lateral views). - Prompt sensitivity: small changes to the prompt template materially change suggestion quality.
- Without
ANTHROPIC_API_KEY, only heuristic dry-run mode is available.
Stack
What this project uses, and why
- Python
- Typer
- Anthropic SDK
- sqlglot
- pytest
See the full code
The repo runs locally with no cloud account — architecture doc, tests, and reproducible measurements included.