Analysis MergeNode optimisations

[davmlaw]

🤖 Written by Claude

Overview

When multiple SampleNode/CohortNode arms feed into a MergeNode, and those arms are over the same cohort (same CohortGenotypeCollection), the merge currently runs each arm as a separate full pass over the cohortgenotype partition, materialises each arm's PKs into a Python list, and combines them as Variant WHERE id IN (listA) OR id IN (listB).

The arms already share one join alias (cohortgenotype_<cgc_pk>), so in principle the same result could be produced in a single pass — exactly what a CohortNode does (one join, one OR/regex predicate over samples_zygosity). The duplicate passes are the optimisation target.

Why it happens

arg_q_dict expresses AND-of-keyed-filters, applied one annotation alias at a time (AnalysisNode.get_queryset). A merge needs OR across arms, and:

• Each sample's zygosity filter is keyed under a per-sample alias (sample_<pk>, a Substr over the shared join), so distinct arms never share a key.
• The format has no representation for "OR across different annotation-dependent keys".

So MergeNode._split_common_filters can only OR-combine None-keyed (annotation-independent) filters; any arm with a non-None key falls back to parent.get_queryset() → materialise PKs → Q(pk__in=...).

Notes / scope

• Applies to ancestors, not just direct parents. arg_q_dict and its stable aliases propagate up through intermediate filter nodes unchanged, so the same-cohort arms are detectable however deep the SampleNode sits below the merge. This is why a fix keyed on shared annotation join generalises, whereas a fix keyed on "are my parents SampleNodes" would miss the (common) case of filtering on each arm before merging.
• Only large arms matter. Arms resolving to ≤ ANALYSIS_NODE_STORE_ID_SIZE_MAX (default 1000) are already collapsed to literal PK lists upstream (get_small_parent_arg_q_dict), so they neither benefit nor need the rewrite. The saving lands on large arms — which is also where the duplicate passes hurt most.
• Gate on shared join, not blanket OR. Combining into one pass is only a win when arms share the same cohortgenotype join. Cross-cohort arms would multiply LEFT JOINs under the OR and likely regress — they should stay on the current materialise-then-OR path (the existing worst-case-protection behaviour).

Possible approaches

1. Special-case same-cohort merge — detect arms over one CohortGenotypeCollection and rebuild a cohort-style single-pass query. Simple-ish but a parallel query path that must re-handle quality/AF/gene-list/VCF/sub-cohort filters, and breaks on intermediate filter nodes if keyed on node class.
2. General filter rewriting (preferred) — extend the merge/get_queryset protocol so an OR-Q can reference multiple already-annotated aliases that share a join: annotate them all up front, then apply a single OR .filter(). Naturally handles ancestors and degrades safely to today's behaviour for distinct joins.

Relevant code

• analysis/models/nodes/filters/merge_node.py — _split_common_filters, _get_merged_q_dict
• analysis/models/nodes/analysis_node.py — get_queryset (annotate-then-filter loop), get_arg_q_dict, get_small_parent_arg_q_dict
• analysis/models/nodes/cohort_mixin.py, snpdb/models/models_vcf.py (Sample.get_cohort_genotype_alias_and_field), snpdb/models/models_cohort.py (cohortgenotype_alias, get_zygosity_q)

[davmlaw]

🤖 Written by Claude

Drafted an implementation plan for the same-cohort MergeNode single-pass optimisation and reviewed it against the current code. Summary below.

Approach

When multiple arms that filter over the same cohort (same CohortGenotypeCollection / cohortgenotype_<cgc> join) feed a MergeNode, produce the merged result in one pass over the partition instead of running each arm as a full standalone pass, materialising its PKs, and OR-ing Variant WHERE id IN (listA) OR id IN (listB).

This is approach 2 (general filter rewriting) from the issue. The key enabler already exists: MergeNode.get_queryset calls get_annotation_kwargs(), which unions all parents' kwargs — de-duping the shared cohortgenotype_<cgc> FilteredRelation while keeping each arm's distinct sample_<pk> Substr. So by the time the None-bucket filters run (applied last, after the whole annotate loop), every arm's alias is already annotated. Placing a single OR-Q in that bucket:

WHERE (SUBSTR(cg.samples_zygosity, idxA, 1) IN ('E','O'))
   OR (SUBSTR(cg.samples_zygosity, idxB, 1) IN ('E','O'))

yields one LEFT JOIN, two Substrs, one OR — no protocol change to get_queryset needed.

Gating (keeps the wins targeted, avoids regressions)

• Shared single join only. Only arms that provably sit over the same cohortgenotype_<cgc> FilteredRelation combine. Cross-cohort / multi-join arms stay on the existing materialise path (OR-ing across different joins would multiply LEFT JOINs and regress).
• Large arms only. Arms ≤ ANALYSIS_NODE_STORE_ID_SIZE_MAX are already collapsed to literal PK lists upstream — untouched.
• Ancestors, not just direct parents. Detection keys on the shared annotation join (which propagates up through intermediate filter nodes), so SampleNode → ZygosityNode → Merge is still detected.
• Group size ≥ 2 to combine; a lone same-cohort arm stays on materialise.
• disable_cache / picklability invariants preserved (the OR-Q embeds only aliases + literals, no lazy queryset — MergeNode - remove cache parents - optimisation with common parents #240, Speed up small jobs by using explicit variant IDs? #546 concerns don't apply).
• Off-switch ANALYSIS_MERGE_NODE_SINGLE_PASS (default on) restores exact current behaviour for rollback and equivalence tests.

Review finding (folded into the plan)

The one substantive gap was the alias→join resolver. CohortNode filter keys don't all take the bare-FilteredRelation-alias shape:

CohortNode config	filter key	naive resolver
COUNT ref/het/hom	cohortgenotype_<cgc>	resolves ✅
COUNT het-or-hom	cohortgenotype_<cgc>__non_ref (F-expr)	drops to materialise ❌
SIMPLE_ZYGOSITY	cohortgenotype_<cgc>__het_count	drops to materialise ❌

A resolver matching only isinstance(v, FilteredRelation) is safe (correct results) but never fires the CohortNode optimisation. The plan now specifies a three-shape resolver (bare alias → sample_<pk> Substr → <cgc_alias>__<suffix> prefix, validated against the arm's FilteredRelation set), and the CohortNode test asserts all three shapes combine.

Verification strategy

Single-pass is proven structurally, not by wall-clock: render the merged SQL and assert the cohortgenotype_<cgc> join appears exactly once (vs once-per-arm / pk IN per arm on the legacy path). Every equivalence test runs twice (flag on/off) and asserts identical PK sets, plus union-of-arms equality. Sub-cohort arms (extra any-sample-called VC join) and cross-cohort arms are tested to stay on materialise.

Plan file: claude/issue_1619_merge_node_optimisation_plan.md.