Let DuckDB write the JSON export, and stop sorting it (~3x faster)

[gaurav]

Summary

Makes the JSON export ~3x faster by deleting the two things the last run's progress log pointed at: a global sort that had to finish before a single row could be written, and a single-threaded Python serialization loop that then did all the work on one core. Closes #8.

What the log showed

From the 2026-08-05 full-corpus run (40,923,261 documents, 18m 06s):

• The first progress line landed at elapsed 3m 00s with exactly 5,000 documents — one fetchmany(batch_size=5000). Nothing had been written for three minutes because ORDER BY la.pmid must materialize and sort all 40.9M rows first. Reproduced locally on 5M rows: first batch at 0.02s without the sort, 2.04s with it.
• RSS sat flat at 161 GiB through the write phase while peak RSS was 201 GiB — the peak belongs to the sort, not the writer.
• Instantaneous throughput decayed 62k → 36k docs/s across the run (the cumulative 39k hides this). Sorted by ascending PMID, later records are newer and far more likely to carry an abstract, so bytes-per-document rise as it goes. Seven of eight cores were idle throughout.

What changed

DuckDB writes the JSON — one COPY (...) TO <dir> (FORMAT JSON, PER_THREAD_OUTPUT true) instead of a fetchmany loop calling json.dumps per row, so serialization runs in C++ on every thread. The query no longer sorts.

Measured end-to-end on a 2M-document copy of the database (same machine, same query, --shards 8):

	Wall time	Rate
Python loop + ORDER BY	112.9s	17.7k docs/s
COPY, no sort	35.5s	56.2k docs/s

The two outputs are identical: 1,995,600 rows each, and EXCEPT in both directions over the full record set returns nothing. (The new files are ~1.5% smaller — DuckDB writes compact JSON where json.dumps defaults to ", " / ": " separators.)

Keeping the record shape honest

_JSON_FIELDS — output field name → SQL expression, in emitted order — is now the single definition of a document, and the COPY projection is built from it. So the DocumentMetadataAPI names, the empty-string-not-null rule and the field order are written down once. validate imports JSON_FIELDS instead of probing _document's arity, which deletes that hack.

month_to_abbrev and _year_from_medline_date stay as Python, because validate normalizes the efetch side with them. Their SQL twins (_PUB_MONTH_SQL, _PUB_YEAR_SQL) are built from the same _MONTH_ABBR / regex and pinned to the Python by two tests over the cases an index-based lookup gets wrong ("0", "13", "99999999999999999999", "Sept", "SEPTEMBER", whitespace, None). Both were mutation-tested: dropping the capitalization from the month key, or loosening the year regex to ([0-9]{4}), fails them.

Three behaviour changes to know about

• --shards N is a maximum, not a count. Output is one file per writer thread, so --shards caps that statement's threads (restored afterwards) and a small dataset can be written by fewer. Default is DuckDB's own thread count. On Slurm, match it to --cpus-per-task.
• Shards are gzipped by default (pubmed_metadata_0.ndjson.gz); --no-gzip opts out. NDJSON compresses ~4-5x, so a full corpus lands at roughly 12 GiB rather than 52 — less written by the export, less read back by validate, less kept around. Compression happens as each shard is written, costing CPU (which this PR just freed up) rather than a second pass. Safe only because nothing downstream needs telling: find_shards already matched both extensions, and check_structure reads through a raw handle so its byte-progress denominator is the compressed size either way. test_cli_export_then_validate_needs_no_flags runs export then validate with no flags on either to keep that true.
• Records are no longer in PMID order. Nothing downstream consumes the order — the ingest is an ElasticSearch bulk load, and validate builds its own sorted PMID manifest.

Also: DuckDB appends to a per-thread output directory rather than clearing it, so a shorter run would have left the previous run's shards behind to be read as current. The export now deletes its own pubmed_metadata_* files first (tested).

Progress output

Per-batch progress lines go with the Python loop. A heartbeat logs output size and current RSS once a minute in their place — no ETA, since the total output size isn't known until it's written:

INFO pubmed2db.export: starting JSON export: 40923261 document(s) to at most 16 shard(s) in data/json
INFO pubmed2db.export: writing: 22.4 GiB across 16 shard(s) · 187 MiB/s · elapsed 2m 03s · RSS 143.1 GiB
INFO pubmed2db.export: exported 40923261 documents to 16 shard(s) in data/json in 5m 12s (peak RSS 88.0 GiB)

(Shapes, not measurements — see below.) -v additionally enables DuckDB's own progress bar.

Verification

• 112 tests pass, 1 skipped (the /proc current-RSS check, which skips on macOS).
• The record-content tests were left untouched and still pass — the SQL projection produces the same documents as _document did, including identifier ordering, the month abbreviation and the MedlineDate year recovery.
• New tests cover the shard cap being a maximum, the thread setting being restored, stale shards from a previous run being removed, the heartbeat line rendering, non-ASCII reaching the file as UTF-8 rather than \uXXXX escapes, and the export→validate round trip with no flags on either command.
• validate's fixture is now a gzipped export, so its whole suite runs against what the CLI actually writes (and exercises appending a second gzip member to a shard).

Not verified

The new peak RSS and wall time on the cluster. The 3x is a laptop measurement on a 2M-document database; the full corpus is 20x that and the machine is different. Two things to record from the first real run:

1. Peak RSS. The sort was the peak-memory event, so --mem=256G is probably now over-provisioned — but by how much is a measurement, and under-requesting is an OOM kill several minutes in. slurm/README.md says to keep 256G until a new number exists.
2. Wall time, which decides whether --time=02:00:00 is still generous.

TODO

• Record the new export's peak RSS and wall time from a cluster run. The 3x is a 2M-document laptop measurement; the full corpus is 20x that on a different machine. slurm/README.md deliberately still says to request 256 GB until a real number exists — this is what replaces it, and it is also the first measurement of the gzipped default's CPU cost.
• Decide what --sample-size means now that shard count is variable. validate samples --sample-size records per shard, so a 16-shard export at the default samples 240 — but with one file per writer thread, the shard count is whatever the export's parallelism happened to be, and the total sample silently moves with it. Either make the flag a total (dividing across shards) or document the coupling; leaving it is the option that quietly changes how much gets checked between runs.

🤖 Generated with Claude Code