[d0choa]

This ports the variant index step from #205 to polars, replacing the Spark variant step on main.

@project-defiant, as the author of the Gentropy migration in #205, your view on this would be very useful.

Where I stand

The variant step sits on the boundary between the two worlds, and both implementations have real pros and cons. I now lean towards polars. It needs about half the compute, runs on the default PTS VM, and leaves Dataproc to Gentropy alone. Since review it also shares its hashing, scaling and part-listing helpers with the rest of PTS, and FoldX runs in the same step. The Spark version is still less code, and it would scale further if the variant index grew by an order of magnitude. Both sides are laid out below.

Spark vs polars

Both versions ran on the same inputs (sz/ppp-2610-1), and their outputs were compared value by value, one chromosome at a time.

	PySpark (#205)	Polars (this PR)
Machines	Dataproc: 1 × n1-standard-8 master + 4 × n1-highmem-16 workers	1 × n2d-standard-32 GCE VM, the default PTS shape
vCPUs / memory	72 vCPUs / 446 GB	32 vCPUs / 128 GB (41 GB peak)
Wall time	~16 min including cluster creation	17.5 min (7 min parsing VEP, 10 min annotating and writing)
vCPU-minutes	~1,150	~560
Step source lines (excluding consequences.py and schema)	~680	~800
Tests	28 tests	39 tests (63 cases), ~8 s, no JVM

Parity: the outputs match.

• Both have 8,484,816 rows, and the schemas are identical on all 25 chromosomes.
• Exactly one value differs. For 7_142349257_G_C, two FoldX predictions tie and each implementation keeps a different one. Neither choice is more correct than the other.
• The comparison ran against feat(pis,pts): generate variant annotations and the variant index within the pipeline #205 before its last changes. Both of those changes have been ported since: OMIM IDs are read as text, and an empty annotation source fails the step.
• The comparison also predates the review changes below, two of which move values on purpose:
  • AlphaMissense, CADD and GERP normalisedScore values are now rescaled in double precision with the shared min_max_scaler, so they can differ from Spark in the last digits.
  • Variant IDs of exactly 300 characters are now hashed in evidence and LoF curation too, as the variant index already did.
• OMIM IDs differ from main on purpose. VEP writes OMIM IDs as JSON numbers and sometimes drops the trailing zero of the four-digit variant suffix. In a sample of 50 VEP parts, 17 of 401 IDs looked like 159530.001 for 159530.0010.
  • Spark on main reads the raw text, so it writes 159530#001 for these.
  • Polars can't read a number's raw text, so it pads the suffix back to four digits and writes 159530#0010, the real OMIM ID.
  • The same padding would fix the Spark version.

Arguments for polars

• No cluster. The step runs on the same GCE VM as most of PTS, at the default machine type, with no override. It drops the pts_variant Dataproc profile, the init action and the Spark settings the step depended on (ANSI off, adaptive query execution, zstd codec). If this goes ahead, Dataproc stays only for Gentropy.
• About half the compute. It uses half the vCPU-time and under a third of the memory for roughly the same wall time.
• Same direction as the rest of PTS. We prefer transformers for new work, and the recent ports (ontoma, evidence post-processing, ...) all moved off Spark. A Spark-only step keeps the PTS Dataproc path alive for one step. It also stays exposed to Spark upgrades, such as Spark 4 reordering collect_set and changing type coercion under ANSI.
• Cheap to develop and test. The tests run in seconds without a JVM, and the step can be run and debugged locally on a slice. Extensions are cheaper to write and to verify.
• Common extensions are small:
  • A new VEP field is one entry in schemas/vep.py plus one expression.
  • A new variant-level or amino-acid-level annotation source is an entry in config.yaml, provided it follows variant_schema or the FoldX shape.
• Room to grow on one machine. Memory is bounded by the largest chromosome (chr1, ~740k variants) plus its slice of gnomAD.
  • n2d-highmem-32 gives about 6× the current peak; pts_evidence_europepmc already uses that override.
  • n2d-standard-64 would speed up parsing, which scales with the number of processes.

Arguments for PySpark

• Less, more declarative code. The polars version adds orchestration that Spark does for us: a loop over chromosomes, batching, a process pool, a scratch directory, and a hand-written union that touches only the rows that need it. That is about 20% more code, all of it specific to this step.

• Horizontal scaling. Polars can only grow by moving to a bigger VM.

  • Growth of 2–3× fits easily.
  • Around 10× needs partitioning by position window rather than by chromosome.
  • An index at gnomAD scale (hundreds of millions of variants) would not fit this design. Spark would only need more workers.

• Nested data is the weak spot of polars 2. The variant index is mostly list[struct] columns, and the port hit several polars traps along the way:

  • a multi-key sort_by inside list.eval fails on empty lists;
  • list.arg_max returns the last maximum for strings but the first for floats;
  • struct when/then panics on multi-chunk frames;
  • replace_strict inside list.eval is rebuilt for every row;
  • eq_missing on nested lists is very slow.

  The tests now cover each of these, but new nested logic is likely to hit more. Spark's array_union and sort_array are mature.

• Gentropy lineage. The step comes from Gentropy's variant_index. If we expect to keep pulling changes from Gentropy, the Spark version can take them almost line for line, while the polars version needs each one re-ported.

Changes since review

• 5d0eaad: annotation parts whose columns differ no longer fail chromosome discovery.
• fd347ea: a shared list_parts helper in utils/dataset.py replaces three copies of the part-listing logic, and the FoldX lookup keys are only computed when a FoldX source is configured.
• f2ccc4b: the step becomes one transformers/variant/ package, laid out like transformers/target/.
• 348054a: score normalisation uses the shared min_max_scaler instead of a Spark-precision _rescale.
• c6168f6: FoldX runs inside the variant step, and pts_foldx leaves the DAG.
• 75d5f93: one variant ID hashing rule for evidence, LoF curation and the variant index.

What this PR contains

Step

• pts/src/pts/transformers/variant/__init__.py is a new variant transformer that replaces the pyspark/variant.py job.
  • VEP JSON parts are parsed in a pool of spawn processes, one per CPU by default, each with one polars thread. Each part is written to its own local Parquet file.
  • The index is then annotated one chromosome at a time. The step reads only that chromosome's gnomAD and LoF curation parts, annotates batches of variants in parallel, and writes each chromosome as one part of the output.
  • Settings:
    • hash_threshold;
    • parse_processes (one per CPU by default) and parse_threads (polars threads per parsing process, default 1);
    • batch_size.
• pts/src/pts/transformers/variant/vep.py parses VEP records:
  • transcript consequences;
  • variant effects (VEP, LOFTEE, AlphaMissense, CADD, GERP);
  • database cross-references, including OMIM and ClinVar;
  • variant descriptions.
• pts/src/pts/transformers/variant/index.py:
  • normalises scores with the shared min_max_scaler;
  • merges annotations with Spark's array_union semantics and its descending sort of struct arrays;
  • adds FoldX annotations through each variant's first transcript with an amino-acid change;
  • fails the step on an empty annotation source, as the Spark version does.
• pts/src/pts/schemas/vep.py holds a VEP schema pruned to the fields the step reads, so that scan_ndjson skips everything else.
• pts/src/pts/transformers/utils/variant.py has one hash_long_variant_ids for evidence, LoF curation and the variant index.
  • IDs of 300 characters or more are hashed. Evidence and LoF curation used to hash only longer ones, so a 300-character evidence variantId could never join the index.
  • Coordinates come from columns when the caller has them, otherwise from the ID, as evidence did.
• pts/src/pts/transformers/utils/dataset.py gains list_parts, which scan_dataset and the variant step share to list a dataset's parts.
• variant/consequences.py moves from pyspark/variant_utils/ unchanged.

Config and orchestration

• The task in pts/config.yaml and ppp/pts.override.yaml becomes transform variant index with transformer: variant, and output_partitions goes. Without a pyspark … task, pts_variant runs on a GCE VM.
• FoldX now runs inside the variant step. transform FoldX annotations is the step's first task and transform variant index requires it, in both pts/config.yaml and the PPP override.
• pts_foldx is removed from unified_pipeline.yaml, and pts_variant now depends on pis_variant and gentropy_variant_annotation.
• The pts_variant Dataproc profile is removed from clusters.yaml.
• The orchestration tests now expect pts_variant to run on GCE, with FoldX inside it.

Removed

• pts/src/pts/pyspark/variant.py
• pts/src/pts/pyspark/variant_utils/ (spark.py, index.py, vep.py)
• pts/tests/test_pyspark_variant.py

Tests

• pts/tests/test_variant.py ports all of the Spark tests.
• New tests cover:
  • OMIM IDs keeping their four-digit suffix;
  • failing on an empty gnomAD or FoldX source;
  • score normalisation;
  • the struct sort order;
  • ties in the highest score, including a LOFTEE tie;
  • repeated values removed only where an annotation was joined;
  • grouping annotation parts by chromosome, including parts whose columns differ;
  • hashing an evidence ID of exactly 300 characters, and rejecting a chromosome passed without a position;
  • the step end to end, and refusing an existing destination.

How it was verified

• Full runs on an n2d-standard-32 VM against sz/ppp-2610-1, compared value by value with the Spark output of feat(pis,pts): generate variant annotations and the variant index within the pipeline #205 for the same run. The results are above.
• The full pts and orchestration test suites pass, as do ruff and ty.

🤖 Generated with Claude Code

https://claude.ai/code/session_015KGXafVcd7q9w8y1uKHc7u

[project-defiant]

Impressive work, if we are going to make the variant lighter (which would make sense due to the fact that it mixes the transcript as well) probably it's worth to get rid of pyspark.

[project-defiant]

Probably worth joining with existing solution for evidence.

[DSuveges]

Oh, I made the same comment, I haven't seen yours.

[project-defiant]

I think this also lives for min_max_scaler for foldx or lof, worth to check if can be reused

[d0choa]

Done in 348054a. _rescale existed only to reproduce Spark's mixed single/double precision, and that isn't worth a second scaler. AlphaMissense, CADD and GERP normalisedScore values may move in the last digits against the previous release.

[project-defiant]

Now that the step is run on a single machine, I would get rid of foldx step and embed it here. No need to spawn machine two times.

[d0choa]

Agreed, done in c6168f6. FoldX is now the first task of the variant step, and pts_foldx is removed from the DAG. Only the variant index reads its output.

[DSuveges]

Shouldn't this method be used in other resources where we need variant hashing? eg. EVA/ClinVar?

[DSuveges]

I also not sure if we should have transformers/utils/variants.py when we have transformers/variant_utils.

[d0choa]

Ironically, the two hashing functions were behaving differently when the variant had exactly 300 characters 🙃. This is another benefit of converging codebases

[DSuveges]

Dang. I did a lot of checks when the duplicated function was added to gentropy and the etl, but I might have missed the equality. Also, @SzymonSzyszkowski pointed out that when the evidence logic was migrated to polars, the hashing function produces different output. In relation to generating evidence ids, I knew this,but forgot the implication this discrepancy caused in the variant index generation.

[d0choa]

On the naming: good point. The step is now one transformers/variant/ package (f2ccc4b), laid out like transformers/target/. transformers/utils/variant.py stays for the helpers that evidence, LoF curation and FoldX share with it.