Skip to content

Commit d998589

Browse files
JuliaEdomclaude
andcommitted
feat(eval): run the pipeline on Postgres against BIRD's own Postgres gold
The "Postgres 16 as a second backend" claim had never been exercised on data: the read-only proof ran against an empty probe table in CI, and every eval was SQLite. This makes the claim real and measures it. BIRD ships a Postgres edition: `MINIDEV_postgresql/BIRD_dev.sql` (955 MB pg_dump of all eleven databases into one `public` schema) and `mini_dev_postgresql.json`, whose gold is rewritten for PG in 312 of 500 questions (STRFTIME → TO_CHAR …). Restoring the dump whole would drop 75 unrelated tables in front of schema retrieval, so `extract_pg_dump_slice.py` cuts one database out of it by table name — globally unique across the eleven, a property a test now pins — and validates the slice (every CREATE TABLE, every COPY, no FK pointing outside). COPY payloads are never scanned for object headers: a StackExchange post body contains lines starting with `--`, and treating one as a boundary would truncate `posts` mid-load, silently. `load_postgres.py` (SQLite → pandas → PG) can't serve this: BIRD's PG gold is written against the dump's schema, whose identifiers are lower-cased (`displayname`), while pandas reproduces SQLite's CamelCase. `SELECT DisplayName FROM users` would not resolve. It stays the right tool for Chinook. Wiring: `--dialect postgresql` (loads PG gold, prompts the model for PG SQL), `--pg-dsn`/`--pg-db-id` on eval_baseline and build_index, and the registry now registers Postgres last so a matching id deliberately re-points that BIRD database at PG. A guard refuses to run PG gold against a SQLite engine — that combination would not error loudly, it would quietly deflate EA. Result (codebase_community, 741 646 rows, config E / codestral, hints off, same 49 questions): Postgres 49.0% (24/49) vs SQLite 44.9% (22/49), validity 100%, 88% agreement. Two questions on n=49 is noise, not superiority — the finding is that the pipeline ports without degrading. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 8baa15a commit d998589

11 files changed

Lines changed: 5824 additions & 23 deletions

‎eval/baselines/postgres_codebase_community_n49.json‎

Lines changed: 1703 additions & 0 deletions
Large diffs are not rendered by default.

‎eval/reports/2026-07-11/E_dense_fewshot_repair-pg.json‎

Lines changed: 1702 additions & 0 deletions
Large diffs are not rendered by default.

‎eval/reports/2026-07-11/E_dense_fewshot_repair-sqlite-cc.json‎

Lines changed: 1744 additions & 0 deletions
Large diffs are not rendered by default.

‎eval/reports/2026-07-11/index.html‎

Lines changed: 102 additions & 0 deletions
Large diffs are not rendered by default.

‎scripts/build_index.py‎

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626

2727
from nl_sql.agent.graph import PipelineConfig
2828
from nl_sql.config import get_settings
29-
from nl_sql.db.registry import get_default_registry
29+
from nl_sql.db.registry import DatabaseRegistry, get_default_registry
3030
from nl_sql.llm.cache import CachingEmbeddingProvider
3131
from nl_sql.llm.providers.base import EmbeddingProvider
3232
from nl_sql.llm.providers.mistral import MistralProvider
@@ -53,8 +53,14 @@ def _runtime_sample_size_default() -> int:
5353
runtime advertises."""
5454

5555

56-
def build_for_db(idx: SchemaIndex, db_id: str, *, sample_size: int = DEFAULT_SAMPLE_SIZE) -> int:
57-
registry = get_default_registry()
56+
def build_for_db(
57+
idx: SchemaIndex,
58+
db_id: str,
59+
*,
60+
sample_size: int = DEFAULT_SAMPLE_SIZE,
61+
registry: DatabaseRegistry | None = None,
62+
) -> int:
63+
registry = registry or get_default_registry()
5864
spec = registry.get(db_id)
5965
print(f"[introspect] {db_id} ({spec.url})")
6066
tables = introspect(spec.make_engine(), sample_size=sample_size)
@@ -98,6 +104,22 @@ def build_parser() -> argparse.ArgumentParser:
98104
action="store_true",
99105
help="Disable diskcache wrapper around the embedding provider.",
100106
)
107+
parser.add_argument(
108+
"--pg-dsn",
109+
default="",
110+
help=(
111+
"Postgres DSN. When set, --pg-db-id is served by Postgres instead of "
112+
"SQLite, and the chunks are introspected from the live Postgres schema "
113+
"(lower-cased identifiers, real timestamptz). Build these into a "
114+
"SEPARATE --persist dir: chunk ids are keyed by db_id, so a Postgres "
115+
"build would otherwise overwrite the SQLite chunks for the same db."
116+
),
117+
)
118+
parser.add_argument(
119+
"--pg-db-id",
120+
default="bird_codebase_community",
121+
help="registry id served by --pg-dsn (default: bird_codebase_community)",
122+
)
101123
return parser
102124

103125

@@ -132,12 +154,12 @@ def main(argv: list[str] | None = None) -> int:
132154
)
133155
idx = SchemaIndex(persist_dir=persist, embedder=embedder, client=client)
134156

135-
registry = get_default_registry()
157+
registry = get_default_registry(pg_dsn=args.pg_dsn, pg_db_id=args.pg_db_id)
136158
targets = registry.ids() if args.db == "all" else [args.db]
137159

138160
total = 0
139161
for db_id in targets:
140-
total += build_for_db(idx, db_id, sample_size=args.sample_size)
162+
total += build_for_db(idx, db_id, sample_size=args.sample_size, registry=registry)
141163
print(f"[summary] indexed {total} chunks across {len(targets)} db(s)")
142164
return 0
143165

‎scripts/eval_baseline.py‎

Lines changed: 48 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,31 @@ def main(argv: list[str] | None = None) -> int:
207207
"rebuild needed. Recommended value: 5."
208208
),
209209
)
210+
parser.add_argument(
211+
"--dialect",
212+
choices=["sqlite", "postgresql"],
213+
default="sqlite",
214+
help=(
215+
"SQL dialect of the run. 'postgresql' loads BIRD's Postgres gold "
216+
"(mini_dev_postgresql.json — 312/500 golds are rewritten for PG, e.g. "
217+
"STRFTIME → TO_CHAR) and prompts the model for Postgres SQL. Requires "
218+
"--pg-dsn, and every sampled example must resolve to a Postgres-backed "
219+
"database (default: sqlite)."
220+
),
221+
)
222+
parser.add_argument(
223+
"--pg-dsn",
224+
default="",
225+
help=(
226+
"Postgres DSN serving --pg-db-id (owner or read-only role). Load it "
227+
"first: scripts/extract_pg_dump_slice.py + psql, see docs/03_eval_methodology.md §14."
228+
),
229+
)
230+
parser.add_argument(
231+
"--pg-db-id",
232+
default="bird_codebase_community",
233+
help="registry id served by --pg-dsn (default: bird_codebase_community)",
234+
)
210235
parser.add_argument(
211236
"--provider",
212237
choices=["mistral", "groq", "github_models", "ollama", "perplexity", "openrouter"],
@@ -228,7 +253,11 @@ def main(argv: list[str] | None = None) -> int:
228253
print("[error] invalid --only-qids: expected comma-separated integers", file=sys.stderr)
229254
return 3
230255

231-
examples = load_bird_mini_dev(Path(args.bird_root))
256+
if args.dialect == "postgresql" and not args.pg_dsn:
257+
print("[error] --dialect postgresql requires --pg-dsn", file=sys.stderr)
258+
return 3
259+
260+
examples = load_bird_mini_dev(Path(args.bird_root), dialect=args.dialect)
232261
if args.db:
233262
examples = [e for e in examples if e.registry_db_id == args.db]
234263
if not examples:
@@ -258,7 +287,7 @@ def main(argv: list[str] | None = None) -> int:
258287
return 3
259288
print(f"[info] loaded {len(examples)} examples → sampled {len(sample)} (seed={args.seed})")
260289

261-
registry = get_default_registry()
290+
registry = get_default_registry(pg_dsn=args.pg_dsn, pg_db_id=args.pg_db_id)
262291
missing = sorted({e.registry_db_id for e in sample} - set(registry.ids()))
263292
if missing:
264293
print(
@@ -268,6 +297,23 @@ def main(argv: list[str] | None = None) -> int:
268297
)
269298
return 4
270299

300+
# Gold SQL is dialect-specific: BIRD's Postgres gold uses TO_CHAR/CAST where
301+
# the SQLite gold uses STRFTIME. Executing it against a SQLite engine would
302+
# not error out loudly enough — it would just fail the query and score a
303+
# miss, quietly reporting a bogus EA. Refuse the run instead.
304+
wrong_engine = sorted(
305+
{e.registry_db_id for e in sample if registry.get(e.registry_db_id).dialect != args.dialect}
306+
)
307+
if wrong_engine:
308+
print(
309+
f"[error] --dialect {args.dialect} but these DBs are not backed by "
310+
f"{args.dialect}: {wrong_engine}\n"
311+
f" load one first (scripts/extract_pg_dump_slice.py) and pass "
312+
f"--pg-dsn/--pg-db-id, or narrow the run with --db {args.pg_db_id}",
313+
file=sys.stderr,
314+
)
315+
return 4
316+
271317
settings = get_settings()
272318
if not settings.mistral_api_key:
273319
print("[error] MISTRAL_API_KEY not set in .env", file=sys.stderr)

‎scripts/extract_pg_dump_slice.py‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
"""Cut one BIRD Mini-Dev database out of the official PostgreSQL dump.
2+
3+
BIRD ships `MINIDEV_postgresql/BIRD_dev.sql`: a ~955 MB `pg_dump` plain file
4+
that loads **all eleven** Mini-Dev databases into a single `public` schema.
5+
Restoring it whole would put 75 unrelated tables (Formula-1 circuits, card
6+
games, hospital labs …) in front of the schema-retrieval step — a different,
7+
much noisier task than the single-database one the product actually serves.
8+
9+
So we slice. Table names are globally unique across the eleven databases
10+
(verified against `dev_tables.json`), which makes a by-table-name slice
11+
unambiguous.
12+
13+
Why not `scripts/load_postgres.py` (SQLite → pandas → Postgres)? Because BIRD's
14+
own Postgres gold SQL is written against **this** dump's schema: identifiers are
15+
folded to lower case (`displayname`, `posthistory`) and dates are real
16+
`timestamptz`, whereas a pandas load reproduces SQLite's CamelCase and text
17+
dates. Gold like `SELECT DisplayName FROM users` resolves to `displayname` in
18+
Postgres and would fail against a pandas-built table. Loading the official dump
19+
is what makes the Postgres eval an apples-to-apples BIRD run rather than a
20+
lookalike. `load_postgres.py` remains the right tool for Chinook and for ad-hoc
21+
SQLite slices where no Postgres gold exists.
22+
23+
Usage:
24+
# what would be extracted (no output file written)
25+
python scripts/extract_pg_dump_slice.py --db codebase_community --list
26+
27+
# write the slice
28+
python scripts/extract_pg_dump_slice.py --db codebase_community \
29+
--out .tmp/codebase_community_pg.sql
30+
"""
31+
32+
from __future__ import annotations
33+
34+
import argparse
35+
import json
36+
import sys
37+
from collections.abc import Iterator
38+
from dataclasses import dataclass, field
39+
from pathlib import Path
40+
from typing import TextIO
41+
42+
from nl_sql.paths import under_root
43+
44+
DEFAULT_DUMP = under_root("data/bird_mini_dev/MINIDEV_postgresql/BIRD_dev.sql")
45+
DEFAULT_TABLES_JSON = under_root("data/bird_mini_dev/MINIDEV/dev_tables.json")
46+
47+
# pg_dump attaches a comment block to every object:
48+
# --
49+
# -- Name: badges badges_pkey; Type: CONSTRAINT; Schema: public; Owner: …
50+
# --
51+
# and for data:
52+
# -- Data for Name: badges; Type: TABLE DATA; Schema: public; Owner: …
53+
_NAME_PREFIX = "-- Name: "
54+
_DATA_PREFIX = "-- Data for Name: "
55+
56+
# Object types whose `Name:` field starts with the table they belong to.
57+
_TABLE_SCOPED = {
58+
"TABLE",
59+
"TABLE DATA",
60+
"CONSTRAINT",
61+
"FK CONSTRAINT",
62+
"INDEX",
63+
"DEFAULT",
64+
"TRIGGER",
65+
}
66+
67+
68+
@dataclass(frozen=True, slots=True)
69+
class Section:
70+
"""One pg_dump object: its `-- Name:` header plus the SQL that follows."""
71+
72+
name: str
73+
type: str
74+
lines: list[str] = field(repr=False)
75+
76+
@property
77+
def table(self) -> str:
78+
"""Table this object belongs to, lower-cased. Empty when not table-scoped."""
79+
if self.type in _TABLE_SCOPED:
80+
return self.name.split()[0].lower()
81+
return ""
82+
83+
84+
def bird_tables(db_id: str, tables_json: Path = DEFAULT_TABLES_JSON) -> set[str]:
85+
"""Lower-cased table names of one BIRD database, read from dev_tables.json."""
86+
entries = json.loads(tables_json.read_text(encoding="utf-8"))
87+
for entry in entries:
88+
if entry["db_id"] == db_id:
89+
return {t.lower() for t in entry["table_names_original"]}
90+
known = sorted(e["db_id"] for e in entries)
91+
raise KeyError(f"unknown BIRD db_id {db_id!r}. Known: {known}")
92+
93+
94+
def _parse_header(line: str) -> tuple[str, str] | None:
95+
"""Return (name, type) from a pg_dump `-- Name:` comment, else None."""
96+
if line.startswith(_DATA_PREFIX):
97+
rest = line[len(_DATA_PREFIX) :]
98+
elif line.startswith(_NAME_PREFIX):
99+
rest = line[len(_NAME_PREFIX) :]
100+
else:
101+
return None
102+
name, _, tail = rest.partition("; Type: ")
103+
if not tail:
104+
return None
105+
obj_type, _, _ = tail.partition(";")
106+
return name.strip(), obj_type.strip()
107+
108+
109+
def iter_sections(fh: TextIO) -> Iterator[Section]:
110+
"""Stream the dump as objects. The leading SET-block is yielded as PREAMBLE.
111+
112+
Data lines inside `COPY … FROM stdin;` are never scanned for headers — a
113+
StackExchange post body legitimately contains lines starting with `--`, and
114+
treating one as an object boundary would truncate the table mid-load.
115+
"""
116+
current = Section(name="", type="PREAMBLE", lines=[])
117+
in_copy = False
118+
119+
for line in fh:
120+
if in_copy:
121+
current.lines.append(line)
122+
if line.rstrip("\n") == r"\.":
123+
in_copy = False
124+
continue
125+
126+
header = _parse_header(line)
127+
if header is not None:
128+
# The bare `--` line above the header already landed in the previous
129+
# section; harmless (pg_dump ignores stray comments), so leave it.
130+
yield current
131+
current = Section(name=header[0], type=header[1], lines=[line])
132+
continue
133+
134+
current.lines.append(line)
135+
if line.startswith("COPY ") and line.rstrip("\n").endswith("FROM stdin;"):
136+
in_copy = True
137+
138+
yield current
139+
140+
141+
def extract(
142+
dump: Path,
143+
tables: set[str],
144+
out: TextIO | None,
145+
) -> list[Section]:
146+
"""Copy the preamble + every object owned by `tables` from `dump` to `out`."""
147+
kept: list[Section] = []
148+
with dump.open("r", encoding="utf-8", newline="") as fh:
149+
for section in iter_sections(fh):
150+
keep = section.type == "PREAMBLE" or section.table in tables
151+
if not keep:
152+
continue
153+
kept.append(section)
154+
if out is not None:
155+
out.writelines(section.lines)
156+
return kept
157+
158+
159+
def _verify(kept: list[Section], tables: set[str]) -> list[str]:
160+
"""Structural checks — a silently short slice is worse than a loud failure."""
161+
problems: list[str] = []
162+
created = {s.table for s in kept if s.type == "TABLE"}
163+
loaded = {s.table for s in kept if s.type == "TABLE DATA"}
164+
if created != tables:
165+
problems.append(f"CREATE TABLE missing for: {sorted(tables - created)}")
166+
if loaded != tables:
167+
problems.append(f"COPY data missing for: {sorted(tables - loaded)}")
168+
# An FK pointing outside the slice cannot restore — it would reference a
169+
# table we deliberately left behind.
170+
for section in kept:
171+
if section.type == "FK CONSTRAINT":
172+
body = "".join(section.lines)
173+
_, _, ref = body.partition("REFERENCES public.")
174+
target = ref.split("(")[0].strip().lower()
175+
if target and target not in tables:
176+
problems.append(f"FK {section.name!r} references {target!r}, outside the slice")
177+
return problems
178+
179+
180+
def main(argv: list[str] | None = None) -> int:
181+
parser = argparse.ArgumentParser(description=__doc__)
182+
parser.add_argument("--db", default="codebase_community", help="BIRD db_id to extract")
183+
parser.add_argument("--dump", type=Path, default=DEFAULT_DUMP, help="path to BIRD_dev.sql")
184+
parser.add_argument("--out", type=Path, default=None, help="output .sql (omit with --list)")
185+
parser.add_argument(
186+
"--list",
187+
action="store_true",
188+
dest="list_only",
189+
help="print the objects that would be extracted and exit",
190+
)
191+
args = parser.parse_args(argv)
192+
193+
if not args.dump.is_file():
194+
print(f"[error] dump not found: {args.dump}", file=sys.stderr)
195+
return 2
196+
if not args.list_only and args.out is None:
197+
print("[error] pass --out PATH (or --list to preview)", file=sys.stderr)
198+
return 3
199+
200+
tables = bird_tables(args.db)
201+
print(f"[info] {args.db}: {len(tables)} tables → {sorted(tables)}")
202+
203+
if args.list_only:
204+
kept = extract(args.dump, tables, out=None)
205+
else:
206+
args.out.parent.mkdir(parents=True, exist_ok=True)
207+
with args.out.open("w", encoding="utf-8", newline="") as out_fh:
208+
kept = extract(args.dump, tables, out=out_fh)
209+
210+
by_type: dict[str, list[str]] = {}
211+
for section in kept:
212+
if section.type == "PREAMBLE":
213+
continue
214+
by_type.setdefault(section.type, []).append(section.name)
215+
for obj_type in sorted(by_type):
216+
print(f" {obj_type}: {', '.join(sorted(by_type[obj_type]))}")
217+
218+
problems = _verify(kept, tables)
219+
if problems:
220+
for problem in problems:
221+
print(f"[error] {problem}", file=sys.stderr)
222+
return 4
223+
224+
if args.out is not None:
225+
size_mb = args.out.stat().st_size / 1e6
226+
print(f"[ok] wrote {args.out} ({size_mb:.0f} MB)")
227+
else:
228+
print("[ok] slice is structurally complete (use --out to write it)")
229+
return 0
230+
231+
232+
if __name__ == "__main__":
233+
sys.exit(main())

0 commit comments

Comments
 (0)