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.

One divergence got through that pinning and is fixed here: SQL trim(x) strips spaces only, where Python's .strip() strips all whitespace. So "\tMar", "3\n" and "\n3" each gave Mar from month_to_abbrev and "" from _PUB_MONTH_SQL. That is reachable, not theoretical — parse._raw_pubdate stores findtext("Month") verbatim, so a <Month> element spanning a line arrives here intact, and the record would then export a blank pub_month and be reported as a mismatch against efetch, which is a CORE_FIELDS error. The test's only whitespace case was " 3 ", spaces alone — exactly what bare trim() already handled. Tab, newline, CR and mixed cases added.

Three behaviour changes to know about

• --shards N is a maximum, not a count — and it is now the write parallelism. 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 it must match --cpus-per-task, so config.sh derives SHARDS from SLURM_CPUS_PER_TASK rather than leaving the two to be kept in step by hand — a fixed 16 against 04-export.sbatch's 8 CPUs would oversubscribe 2:1 on the one step that gets OOM-killed. test_shards_tracks_the_export_allocation reads --cpus-per-task out of the sbatch header and fails if the fallback disagrees with it.
• Shard names lost their zero padding. pubmed_metadata_00000.ndjson is now pubmed_metadata_0.ndjson.gz: PER_THREAD_OUTPUT names files from FILENAME_PATTERN '{i}', which emits a bare index. Anything globbing pubmed_metadata_* is fine; anything matching the padded form, or relying on the files sorting lexically past nine, is not.
• validate --sample-size is per shard, so the sample now moves with the allocation. It always was per shard, but the count was fixed at 16; with one shard per writer thread, halving the export's --cpus-per-task halves how many records the next validate compares against Entrez. Documented in --help and slurm/README.md rather than changed — making the flag a total means reworking check_structure's per-shard reservoir sampling, which is validate's contract rather than this branch's subject. The report already names what a run checked (240 records sampled: 15/shard x 16 shards).
• 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.

The heartbeat is the only thing standing between an operator and twenty silent minutes, so it is written not to fail: it stats shards that DuckDB is concurrently writing, skipping any that vanish between the glob and the stat, and the loop logs and continues on anything else rather than letting an exception kill the thread. It also starts after the statement that reads the connection's thread count — started before it, a raise there would leave a non-daemon thread looping forever and hang the CLI instead of reporting the error.

The exported ... in <duration> figure covers the whole export, including materializing _latest_snapshot — a window function over the entire article table, and plausibly the dominant phase now the sort is gone. slurm/README.md tells operators to size --time from that line, so a clock started after it would understate the job.

Verification

• 198 tests pass, 1 skipped (the /proc current-RSS check, which skips on macOS), on both CI Pythons — 3.11 and 3.14. Main was merged in partway through, which is where the tests workflow, the Slurm pipeline and the DuckDB --memory-limit work arrived from (Add a Slurm sbatch pipeline, narrate the long-running steps, and cap DuckDB with --memory-limit #9).
• 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 04-export.sbatch's --mem is probably now over-provisioned — but by how much is a measurement, and under-requesting is an OOM kill several minutes in.
2. Wall time, which decides whether the header's --time is still generous.
3. The gzip default's CPU cost, which has never been measured at all — it happens on cores this PR freed up, but that is an argument, not a number.

Tracked in #42. The #SBATCH headers are what to edit once the numbers exist; slurm/README.md carries the run history and deliberately does not restate the allocation.

What the merge from main changed

main gained the Slurm pipeline and CI (#9) while this branch was open, and the merge was not purely mechanical. Three things a reader would otherwise have to reverse-engineer from the merge commit:

• The orientation file moved. main split CLAUDE.md into AGENTS.md; this branch had kept editing CLAUDE.md. Git flagged no conflict — AGENTS.md arrived as a clean new file — so this branch's four decisions (the COPY writer, the gzip default, the absent ORDER BY, EXPECTED_FIELDS from JSON_FIELDS) would have vanished silently. They are ported into AGENTS.md.
• ID_PREFIXES now says PMCID, not PMC. That is main's deliberate bet from Settle the CURIE prefix for PMCIDs in the JSON export's identifiers #33, which this branch predated. It changes every PMCID CURIE in the export, and it arrived through the merge rather than through either PR's own work — worth knowing when reading the identifier diff. The CASE is derived from ID_PREFIXES rather than hand-written, which is pinned by a test.
• Four of main's tests asserted the pre-COPY shape — zero-padded shard names and plain-.ndjson reads that the gzip default breaks. Rewritten against the properties, using the gzip-aware helpers this branch already had.

Follow-on work

• Re-baseline the JSON export's peak RSS and wall time after the COPY rewrite #42 — re-baseline the export's --mem and --time from a cluster run, and measure the gzip default's CPU cost for the first time. The only item here that cannot be settled without hardware.
• FUTURE.md, "Confirm the Node Annotator JSON contract" — upgraded from eventually to before the next production export. Two artifact changes are invisible to the downstream ElasticSearch ingest, which lives outside this repo: the gzip default renames every shard, and an ingest that globs *.ndjson without decompressing reads nothing — an empty ingest rather than an error, which is the worst shape a failure can take. --no-gzip restores the old artifact if the answer turns out to be no.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Reworks JSON export for faster parallel DuckDB serialization while removing the global PMID sort.

Changes:

• Uses DuckDB COPY with per-thread output.
• Makes gzip default and shard count a thread cap.
• Updates validation, tests, and documentation.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
src/pubmed2db/export.py	Implements parallel JSON export and progress logging.
src/pubmed2db/cli.py	Updates export options and defaults.
src/pubmed2db/validate.py	Derives expected fields from the exporter.
tests/test_export.py	Tests export behavior and SQL normalization.
tests/test_cli.py	Tests new CLI defaults and round trip.
tests/test_validate.py	Exercises validation against gzip shards.
README.md	Documents JSON export behavior.
slurm/README.md	Updates performance and operational guidance.
FUTURE.md	Records completion of export optimization.
AGENTS.md	Captures architectural decisions.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.