Julenmendieta/MILAB-6501_handleHugeInputs

[julenmendieta]

Greptile Summary

This PR replaces the memory-bound full-SVD PCA + exact-vector deduplication + sklearn HDBSCAN pipeline with a constant-memory streaming architecture: two-pass IncrementalPCA over the parquet stream, contrib hdbscan (dual-tree Boruvka MST, ~4× faster), and an identity dedup_mapping.tsv (dedup dropped as np.unique over N×D cannot run at scale). A third stream pass populates a refinement store (RAM or disk memmap above 16 GiB) so recursive re-clustering still operates on the original D-dimensional vectors without holding the full matrix.

Key touched terms and their changes:

• IncrementalPCA — sklearn's online PCA. Replaces full-SVD PCA for global reduction; fits over 4M-row batches (pass 1) and transforms in a second pass. Peak memory bounded to N×k (reduced) rather than N×D (raw).
• HDBSCAN (contrib hdbscan 0.8.44) — Replaces sklearn.cluster.HDBSCAN; enables the dual-tree Boruvka MST algorithm on low-dimensional post-PCA spaces (~4× faster). Cluster assignments shift slightly vs sklearn (both are valid HDBSCAN*).
• dedup_mapping.tsv — Was a many-to-one mapping (identical-vector deduplication). Now an identity mapping (representativeKey == clonotypeKey). np.unique over N×D cannot run at scale; kept for schema compatibility in process_results.py.
• Medoid — Actual data-point minimizing the probability-weighted sum of Euclidean distances to cluster members. Computed in the globally-reduced/L2-normalized space. Exact for ≤ 4 000 members, approximate (weighted-mean direction) above that.
• refined_mask / build_refined_store — Boolean mask over N clonotypes identifying points refinement can touch (noise pile + oversized MAIN cluster members). A 3rd stream pass stores only those original D-dimensional vectors, matched by clonotype key, in RAM or disk memmap.
• RAM_BUDGET_GIB (16 GiB) — Threshold controlling whether the refinement store lives in RAM (full-SVD re-PCA) or on a disk memmap (chunked IncrementalPCA). At 10 M × 1024-dim the memmap reaches ~41 GB of scratch disk.
• split_recursive — Recursive closure within run_clustering that re-PCAs a subset (from the refined store), re-clusters with HDBSCAN, assigns fresh global IDs via next_id[0], and recurses into oversized children up to max_depth = 3.
• stream_reduce — Two-pass streaming reduction: pass 1 fits IncrementalPCA over buffered chunks (≥ ncomp rows each); pass 2 re-streams to transform and fill a pre-allocated N×k array.

Confidence Score: 4/5

Safe to merge; the core streaming logic is correct, boundary conditions are well-guarded, and the memmap cleanup uses a try/finally.

The streaming three-pass architecture is correct, the batch-boundary clonotype-split logic handles leftover rows properly, validation in the parquet assembler is thorough (bincount + NaN check + dim-range guard), and the memmap is cleaned up in a finally block even on error. The D naming collision in weighted_medoid and the absent fitted-guard in _reduce_subset's IncrementalPCA path are both practically unreachable under current defaults, but leave the code without defensive safeguards.

software/src/embedding_clustering.py around weighted_medoid and _reduce_subset; software/src/process_results.py line 46 (Python set construction) may matter at very large scale.

Important Files Changed

Filename	Overview
software/src/embedding_clustering.py	Core clustering script fully rewritten: replaces full-SVD PCA + sklearn HDBSCAN + np.unique dedup with streaming IncrementalPCA, contrib hdbscan (Boruvka MST), and identity dedup_mapping. Thorough validation in the streaming loader; refinement store with memmap fallback is correct.
software/src/process_results.py	Adapted for identity dedup_mapping; join logic and downstream aggregation are correct. Python set used for embedded_keys membership check is memory-inefficient for large N but not a correctness issue.
workflow/src/embedding-clustering.tpl.tengo	Memory formula updated to flat formula (parquet_size + 24 GiB, clamped 32-64 GiB) reflecting peak RAM no longer proportional to N×D. CPU API updated to workflow-tengo 6.8.0 onCPU block.
software/src/requirements.txt	Added pyarrow==24.0.0 (streaming loader) and hdbscan==0.8.44 (contrib clustering). Missing trailing newline.
pnpm-workspace.yaml	Bumps all SDK and platform packages to latest versions; runenv-python-3 raised to ^1.11.3 (required by hdbscan). Quote-style normalized to double quotes (cosmetic).

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant P as Parquet Input
    participant SR as stream_reduce
    participant IPCA as IncrementalPCA
    participant BRS as build_refined_store
    participant RC as run_clustering
    participant HDBSCAN as hdbscan.HDBSCAN
    participant WO as write_outputs

    P->>SR: Pass 1 - _stream_clonotypes (batch 4M rows)
    SR->>IPCA: partial_fit(chunk) per batch
    IPCA-->>SR: PCA basis
    SR->>SR: "k = pick_k_95(explained_variance_ratio, 95%)"
    P->>SR: Pass 2 - _stream_clonotypes (transform)
    SR->>IPCA: transform(chunk)[:, :k]
    SR-->>RC: Xr (N x k float32), keys (N,)
    RC->>RC: L2-normalize Xr to Xn
    RC->>HDBSCAN: fit(Xn[valid])
    HDBSCAN-->>RC: labels, probabilities_
    RC->>RC: compute oversized_main + refined_mask
    P->>BRS: Pass 3 - _stream_clonotypes (refined_mask only)
    BRS-->>RC: store (RAM array or disk memmap), gpos
    RC->>RC: split_recursive(oversized MAIN clusters)
    RC->>RC: split_recursive(noise rescue) if rescue-noise
    Note over RC: memmap freed in finally block
    RC->>RC: "cluster_medoids(Xn, labels, weights=probs)"
    RC-->>WO: rep_keys, cluster_id, distance
    WO->>WO: clusters.tsv (headerless)
    WO->>WO: dedup_mapping.tsv (identity mapping)
    WO->>WO: centroid_distances.tsv

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant P as Parquet Input
    participant SR as stream_reduce
    participant IPCA as IncrementalPCA
    participant BRS as build_refined_store
    participant RC as run_clustering
    participant HDBSCAN as hdbscan.HDBSCAN
    participant WO as write_outputs

    P->>SR: Pass 1 - _stream_clonotypes (batch 4M rows)
    SR->>IPCA: partial_fit(chunk) per batch
    IPCA-->>SR: PCA basis
    SR->>SR: "k = pick_k_95(explained_variance_ratio, 95%)"
    P->>SR: Pass 2 - _stream_clonotypes (transform)
    SR->>IPCA: transform(chunk)[:, :k]
    SR-->>RC: Xr (N x k float32), keys (N,)
    RC->>RC: L2-normalize Xr to Xn
    RC->>HDBSCAN: fit(Xn[valid])
    HDBSCAN-->>RC: labels, probabilities_
    RC->>RC: compute oversized_main + refined_mask
    P->>BRS: Pass 3 - _stream_clonotypes (refined_mask only)
    BRS-->>RC: store (RAM array or disk memmap), gpos
    RC->>RC: split_recursive(oversized MAIN clusters)
    RC->>RC: split_recursive(noise rescue) if rescue-noise
    Note over RC: memmap freed in finally block
    RC->>RC: "cluster_medoids(Xn, labels, weights=probs)"
    RC-->>WO: rep_keys, cluster_id, distance
    WO->>WO: clusters.tsv (headerless)
    WO->>WO: dedup_mapping.tsv (identity mapping)
    WO->>WO: centroid_distances.tsv

Loading

Prompt To Fix All With AI

Fix the following 4 code review issues. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 4
software/src/embedding_clustering.py:87-89
**Variable name `D` shadows embedding-dimension convention**

The local `D` assigned here (a pairwise distance matrix, shape `m×m`) reuses the same name that everywhere else in this file denotes the embedding dimension (scalar integer). A reader scanning the function sees `D * w` and must pause to recall that this `D` is a matrix, not the scalar. Since this function is self-contained and doesn't use the outer `D`, there is no runtime bug, but it adds cognitive load whenever the function is edited alongside the rest of the module.

### Issue 2 of 4
software/src/embedding_clustering.py:307-316
**No fitted-guard before accessing `explained_variance_ratio_` in the IncrementalPCA path**

After the chunk loop, if every chunk's row-count is below `ncomp`, `ipca.partial_fit` is never called and `ipca.explained_variance_ratio_` on the next line raises `sklearn.exceptions.NotFittedError`. In practice this path (`not _full_svd_fits`) is only reached for very large subsets (> ~1 M vectors for D = 1024), and `chunk = 200_000 >> ncomp ≤ 500`, so the first chunk always fires `partial_fit`. But the guard is absent — a future change to `chunk` or `ncomp` defaults could expose it silently.

### Issue 3 of 4
software/src/requirements.txt:4-6
The file ends without a trailing newline, which is a POSIX convention violation and can cause diff noise or tool warnings.

```suggestion
numpy==2.2.6
pyarrow==24.0.0
hdbscan==0.8.44
```

### Issue 4 of 4
software/src/process_results.py:46-47
**Python `set` materialisation is O(N) heap for large key sets**

`dedup_mapping.get_column("clonotypeKey").to_list()` materialises all N clonotype keys into a Python list, then `set(...)` copies them into a Python hash-set — two O(N) Python-heap allocations. At 3.3 M clonotypes with ~20-byte keys this is ~130 MB of Python objects before `n_excluded` is even computed. The same exclusion count can be derived with a Polars anti-join: `cloneTable.join(dedup_mapping.select("clonotypeKey"), on="clonotypeKey", how="anti").height`.

_{Reviews (1): Last reviewed commit: "Changeset" | Re-trigger Greptile}

Greptile also left 2 inline comments on this PR.

Context used:

• Context used - Terms is a types in codebase. Provide the list of ... (source)

[gemini-code-assist]

Code Review

This pull request optimizes the embedding-clustering workflow to handle large inputs by switching to a streaming approach with IncrementalPCA, which avoids loading the full N x D matrix into RAM. It also replaces the scikit-learn HDBSCAN implementation with the contrib hdbscan package for improved performance and updates various dependencies. A high-severity issue was identified regarding an invalid pyarrow version in requirements.txt that needs to be corrected.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[PaulNewling]

LGTM! Thanks for the walkthrough