From 847c55f47f5cd1d817be8b7f299a63c03116380c Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 00:22:22 +0200 Subject: [PATCH 01/31] feat(bundle): add parquetbundle v3 encoder Columnar codes, wide float32 projections and a CSR payload part. --- .../src/protspace/data/io/bundle_v3.py | 484 ++++++++++++++++++ apps/protspace/tests/test_bundle_v3_encode.py | 416 +++++++++++++++ 2 files changed, 900 insertions(+) create mode 100644 apps/protspace/src/protspace/data/io/bundle_v3.py create mode 100644 apps/protspace/tests/test_bundle_v3_encode.py diff --git a/apps/protspace/src/protspace/data/io/bundle_v3.py b/apps/protspace/src/protspace/data/io/bundle_v3.py new file mode 100644 index 00000000..45bc2735 --- /dev/null +++ b/apps/protspace/src/protspace/data/io/bundle_v3.py @@ -0,0 +1,484 @@ +"""ParquetBundle format v3: columnar annotation encoding. + +v2 stringifies every annotation cell and packs multi-values as ``;``-joined +hits with ``|``-suffixed scores/evidence, which forces the browser to re-split +and dictionary-code 573K strings on load. v3 moves that work to write time: +part 1 carries int32 dictionary codes (or CSR end offsets) and float64 +numerics, part 3 carries wide float32 projections, and a new part 6 carries the +label dictionaries plus the CSR code/score/evidence payloads as raw +little-endian buffers. + +Only the *container* changes. ``encode_v3`` takes the v2-shaped tables the +pipeline already builds and the (sibling) ``decode_v3`` turns v3 parts back +into them, so every Python consumer keeps its string-cell logic and +``BUNDLE_FORMAT_VERSION = 2`` in :mod:`~protspace.data.annotations.encoding` +still versions the cell grammar. + +The classification rules below intentionally mirror the browser's v2 reader +(``packages/core/src/components/data-loader/utils/conversion.ts``) so a v3 +bundle and its v2 equivalent produce identical colours, code order and legend +entries. Deviations are documented on the constants they come from. +""" + +from __future__ import annotations + +import io +import json +from typing import Any + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq + +from protspace.data.annotations.encoding import ( + FORMAT_VERSION_KEY, + decode_field, + migrate_legacy_annotation_table, + read_format_version, +) + +CONTAINER_VERSION = 3 +MANIFEST_KEY = b"protspace_v3_manifest" + +#: Cell/hit spellings that mean "missing". Mirrors ``MISSING_VALUE_TOKENS`` +#: in ``packages/utils/src/visualization/missing-values.ts``; compared against +#: the lower-cased, whitespace-trimmed token. +MISSING_TOKENS = frozenset({"na", "n/a", "nan", "null", "none", "__na__"}) + +#: ``EVIDENCE_CODE_RE`` from ``conversion.ts``: the part after a hit's last +#: ``|`` is an evidence code, not a score. +EVIDENCE_RE = r"^(?:[A-Z]{2,5}|ECO:\d+)$" + +#: What JavaScript's ``Number()`` accepts *and* ``Number.isFinite`` keeps, +#: restricted to decimal literals. Deviation from the browser: JS also parses +#: ``0x10``/``0o17``/``0b1`` as numbers, so a column of hex literals is +#: categorical here and numeric there. Non-decimal literals do not occur in +#: annotation data and supporting them would cost a Python-level parse. +#: ``Infinity``/``1e999`` are excluded by the post-cast finiteness check. +JS_NUMBER_RE = r"^[+-]?(?:\d+\.?\d*|\.\d+)(?:[eE][+-]?\d+)?$" + +#: hyparquet only hands back zero-copy typed arrays for REQUIRED flat PLAIN +#: columns, so every v3 column is written non-nullable, undictionaried and in +#: one row group. +_PQ: dict[str, Any] = { + "use_dictionary": False, + "column_encoding": "PLAIN", + "compression": "snappy", + "write_statistics": False, +} + +_EVIDENCE_DICT_NAME = "__evidence" + + +# --------------------------------------------------------------------------- # +# encoder +# --------------------------------------------------------------------------- # + + +def _write(table: pa.Table) -> bytes: + """Serialize one v3 part: single row group, PLAIN, no dictionary.""" + buf = io.BytesIO() + pq.write_table(table, buf, row_group_size=max(table.num_rows, 1), **_PQ) + return buf.getvalue() + + +def _required_table( + columns: dict[str, pa.Array], metadata: dict | None = None +) -> pa.Table: + """Build a table whose every field is non-nullable.""" + schema = pa.schema( + [pa.field(name, arr.type, nullable=False) for name, arr in columns.items()], + metadata=metadata, + ) + return pa.table(list(columns.values()), schema=schema) + + +def _as_string(column: pa.ChunkedArray | pa.Array) -> pa.Array: + """Flatten to a single ``string`` array, rendering bools as ``True``/``False``.""" + arr = column.combine_chunks() if isinstance(column, pa.ChunkedArray) else column + if isinstance(arr, pa.ChunkedArray): # combine_chunks keeps the wrapper + arr = arr.combine_chunks() + if pa.types.is_boolean(arr.type): + return pc.if_else(arr, pa.scalar("True"), pa.scalar("False")) + if pa.types.is_string(arr.type): + return arr + return pc.cast(arr, pa.string()) + + +def _missing_mask(trimmed: pa.Array) -> np.ndarray: + """``normalizeMissingValue``: null, blank, or a MISSING_TOKENS spelling.""" + is_null = pc.is_null(trimmed) + blank = pc.equal(trimmed, pa.scalar("")) + token = pc.is_in(pc.utf8_lower(trimmed), value_set=pa.array(sorted(MISSING_TOKENS))) + mask = pc.or_(pc.or_(is_null, blank), pc.fill_null(token, False)) + return np.asarray(pc.fill_null(mask, True)) + + +def _regex_ok(values: pa.Array, pattern: str) -> np.ndarray: + return np.asarray(pc.fill_null(pc.match_substring_regex(values, pattern), False)) + + +def _parse_floats(values: pa.Array, ok: np.ndarray, blank_is_zero: bool) -> np.ndarray: + """Cast the entries flagged by ``ok`` to float64; substitute 0 elsewhere. + + ``Number("")`` is ``0`` in JavaScript, which is how an empty score part + (``"label|1,"``) becomes a real score. + """ + fill = pa.scalar("0") + safe = pc.if_else(pa.array(ok), values, fill) + if blank_is_zero: + safe = pc.if_else(pc.equal(safe, pa.scalar("")), fill, safe) + return pc.cast(safe, pa.float64()).to_numpy(zero_copy_only=False) + + +def _frequency_order(codes: np.ndarray, n_labels: int) -> tuple[np.ndarray, np.ndarray]: + """Return ``(rank, order)`` for the browser's descending-frequency sort. + + ``conversion.ts:1600-1605`` sorts ``Map.keys()`` (first-occurrence order) + with a stable descending-count comparator, so ties keep first occurrence. + """ + counts = np.bincount(codes, minlength=n_labels) + order = np.argsort(-counts, kind="stable") + rank = np.empty(n_labels, dtype=np.int32) + rank[order] = np.arange(n_labels, dtype=np.int32) + return rank, order + + +def _dict_payloads(name: str, labels: list[str]) -> list[tuple[str, bytes]]: + """``dict:`` utf8 blob + ``dict::end`` int32 byte offsets.""" + encoded = [label.encode("utf-8") for label in labels] + ends = np.cumsum([len(b) for b in encoded], dtype=np.int64).astype(" tuple[pa.Array, pa.Array]: + """Split each hit on its LAST ``|`` (``conversion.ts:440``). + + Returns ``(head, suffix_raw)``. A hit without a ``|`` gets ``suffix_raw`` + ``""``, which is the same branch as a trailing-pipe hit: both keep the whole + hit as the label. + """ + parts = pc.split_pattern(hits, "|", max_splits=1, reverse=True) + lengths = np.asarray(pc.list_value_length(parts)) + starts = np.concatenate(([0], np.cumsum(lengths, dtype=np.int64)[:-1])) + flat = pc.list_flatten(parts) + head = flat.take(pa.array(starts)) + has_two = lengths == 2 + suffix = pc.if_else( + pa.array(has_two), + flat.take(pa.array(np.where(has_two, starts + 1, starts))), + pa.scalar(""), + ) + return head, suffix + + +def _encode_annotation_column( + column: pa.ChunkedArray | pa.Array, + name: str, + num_rows: int, + evidence_dict: dict[str, int], +) -> tuple[dict[str, Any], pa.Array, list[tuple[str, bytes]]]: + """Encode one annotation column. + + Returns ``(manifest_entry, part1_array, payloads)``. ``part1_array`` is the + ```` codes / values or the ``__end`` CSR offsets; the caller picks + the physical column name from ``manifest_entry["kind"]``. + """ + source_type = str(column.type) + arr = column.combine_chunks() if isinstance(column, pa.ChunkedArray) else column + + # Arrow-numeric source columns stay numeric regardless of content. The + # browser would call an all-null column categorical, but keeping the kind + # tied to the Arrow type is what lets `decode_v3` restore `sourceType`. + if pa.types.is_integer(arr.type) or pa.types.is_floating(arr.type): + values = pc.cast(arr, pa.float64()).to_numpy(zero_copy_only=False) + values = np.where(np.isfinite(values), values, np.nan) + finite = values[~np.isnan(values)] + numeric_type = "int" if np.all(np.mod(finite, 1) == 0) else "float" + entry = { + "kind": "numeric", + "numericType": numeric_type, + "sourceType": source_type, + } + return entry, pa.array(values, type=pa.float64()), [] + + strings = _as_string(arr) + trimmed = pc.utf8_trim_whitespace(strings) + missing = _missing_mask(trimmed) + + # --- numeric inference (conversion.ts:71-125) --------------------------- # + if not missing.all(): + numeric_ok = _regex_ok(trimmed, JS_NUMBER_RE) | missing + if numeric_ok.all(): + values = _parse_floats(trimmed, ~missing, blank_is_zero=False) + if np.isfinite(values[~missing]).all(): + values = np.where(missing, np.nan, values) + finite = values[~missing] + numeric_type = "int" if np.all(np.mod(finite, 1) == 0) else "float" + entry = { + "kind": "numeric", + "numericType": numeric_type, + "sourceType": source_type, + } + return entry, pa.array(values, type=pa.float64()), [] + + # --- categorical: split cells into hits --------------------------------- # + cells = pc.if_else(pa.array(~missing), trimmed, pa.scalar(None, pa.string())) + hit_lists = pc.split_pattern(cells, ";") + row_of_hit = np.asarray(pc.list_parent_indices(hit_lists)) + hits = pc.utf8_trim_whitespace(pc.list_flatten(hit_lists)) + keep = ~_missing_mask(hits) + if not keep.all(): + hits = hits.filter(pa.array(keep)) + row_of_hit = row_of_hit[keep] + + n_hits = len(hits) + per_row = ( + np.bincount(row_of_hit, minlength=num_rows) + if n_hits + else np.zeros(num_rows, int) + ) + max_hits = int(per_row.max()) if num_rows else 0 + + if n_hits == 0: + entry = {"kind": "categorical", "sourceType": source_type} + codes = np.full(num_rows, -1, dtype=np.int32) + return entry, pa.array(codes, type=pa.int32()), _dict_payloads(name, []) + + # --- per-hit label / score / evidence (conversion.ts:433-468) ----------- # + head, suffix_raw = _split_last_pipe(hits) + no_suffix = np.asarray(pc.equal(suffix_raw, pa.scalar(""))) + suffix = pc.utf8_trim_whitespace(suffix_raw) + is_evidence = ~no_suffix & _regex_ok(suffix, EVIDENCE_RE) + + scored = np.zeros(n_hits, dtype=bool) + hit_score_count = np.zeros(n_hits, dtype=np.int64) + score_values = np.zeros(0, dtype=np.float64) + + candidate = np.flatnonzero(~no_suffix & ~is_evidence) + if candidate.size: + pieces = pc.split_pattern(suffix.take(pa.array(candidate)), ",") + piece_len = np.asarray(pc.list_value_length(pieces)).astype(np.int64) + flat = pc.utf8_trim_whitespace(pc.list_flatten(pieces)) + blank = np.asarray(pc.equal(flat, pa.scalar(""))) + numeric = _regex_ok(flat, JS_NUMBER_RE) + parsed = _parse_floats(flat, numeric, blank_is_zero=True) + valid = blank | (numeric & np.isfinite(parsed)) + owner = np.repeat(np.arange(candidate.size), piece_len) + bad = np.bincount(owner, weights=~valid, minlength=candidate.size) + ok = bad == 0 + scored[candidate[ok]] = True + hit_score_count[candidate[ok]] = piece_len[ok] + score_values = parsed[np.repeat(ok, piece_len)] + + use_head = pa.array(is_evidence | scored) + labels = pc.if_else(use_head, pc.utf8_trim_whitespace(head), hits) + + # --- dictionary in decoded space --------------------------------------- # + encoded_dict = pc.dictionary_encode(labels) + raw_labels = encoded_dict.dictionary.to_pylist() + unify: dict[str, int] = {} + fold = np.empty(len(raw_labels), dtype=np.int32) + for i, raw in enumerate(raw_labels): + fold[i] = unify.setdefault(decode_field(raw), len(unify)) + provisional = fold[np.asarray(encoded_dict.indices)] + rank, order = _frequency_order(provisional, len(unify)) + codes = rank[provisional].astype(np.int32) + ordered_labels = list(unify) + ordered_labels = [ordered_labels[i] for i in order] + + payloads = _dict_payloads(name, ordered_labels) + has_scores = bool(scored.any()) + has_evidence = bool(is_evidence.any()) + + if max_hits <= 1 and not has_scores and not has_evidence: + row_codes = np.full(num_rows, -1, dtype=np.int32) + row_codes[row_of_hit] = codes + entry = {"kind": "categorical", "sourceType": source_type} + return entry, pa.array(row_codes, type=pa.int32()), payloads + + end = np.cumsum(per_row, dtype=np.int64).astype(" tuple[pa.Table, list[dict[str, Any]]]: + """Pivot the long projections table to wide float32, aligned to part 1.""" + required = {"projection_name", "identifier", "x", "y"} + missing = required - set(projections_data.column_names) + if missing: + raise ValueError( + f"projections_data is missing required column(s): {sorted(missing)}" + ) + + names = projections_metadata.column("projection_name").to_pylist() + if len(set(names)) != len(names): + raise ValueError( + f"Duplicate projection name(s) in projections_metadata: {names}" + ) + + dimensions = ( + projections_metadata.column("dimensions").to_pylist() + if "dimensions" in projections_metadata.column_names + else [None] * len(names) + ) + has_z = "z" in projections_data.column_names + + index = pd.Index(protein_ids.to_pylist()) + num_rows = len(index) + columns: dict[str, pa.Array] = {} + manifest: list[dict[str, Any]] = [] + + name_column = projections_data.column("projection_name") + for name, declared in zip(names, dimensions, strict=True): + sub = projections_data.filter(pc.equal(name_column, pa.scalar(name))) + positions = index.get_indexer(sub.column("identifier").to_pylist()) + if len(positions) and positions.min() < 0: + unknown = np.asarray(sub.column("identifier").to_pylist())[positions < 0] + raise ValueError( + f"projection '{name}' references identifier(s) absent from the " + f"annotations table: {sorted(set(unknown.tolist()))[:5]}" + ) + + z = sub.column("z") if has_z else None + z_present = ( + z is not None and not pa.types.is_null(z.type) and z.null_count < len(z) + ) + dimension = int(declared) if declared in (2, 3) else (3 if z_present else 2) + + for axis in ("x", "y", "z")[:dimension]: + values = np.full(num_rows, np.nan, dtype=np.float32) + if axis == "z" and not z_present: + source = None + else: + source = ( + sub.column(axis).to_numpy(zero_copy_only=False).astype(np.float32) + ) + if source is not None: + values[positions] = source + columns[f"{name}__{axis}"] = pa.array(values, type=pa.float32()) + manifest.append({"name": name, "dimension": dimension}) + + return _required_table(columns), manifest + + +def encode_v3( + annotations: pa.Table, + projections_metadata: pa.Table, + projections_data: pa.Table, +) -> tuple[bytes, bytes, bytes, bytes]: + """Encode the v2-shaped pipeline tables as v3 parts 1, 2, 3 and 6.""" + if read_format_version(annotations) == 1: + annotations = migrate_legacy_annotation_table(annotations) + + id_column = next( + (c for c in ("protein_id", "identifier") if c in annotations.column_names), None + ) + if id_column is None: + raise ValueError( + "annotations table has no 'protein_id' or 'identifier' column; " + f"found {annotations.column_names}" + ) + + ids = _as_string(annotations.column(id_column)) + if ids.null_count: + raise ValueError(f"annotations column '{id_column}' contains null values") + duplicated = pc.sum(pc.greater(pc.value_counts(ids).field("counts"), 1)).as_py() + if duplicated: + raise ValueError( + f"annotations column '{id_column}' contains {duplicated} duplicated value(s); " + "protein identifiers must be unique" + ) + + num_rows = annotations.num_rows + existing = set(annotations.column_names) + evidence_dict: dict[str, int] = {} + columns: dict[str, pa.Array] = {id_column: ids} + manifest_columns: dict[str, Any] = {} + payloads: list[tuple[str, bytes]] = [] + + for name in annotations.column_names: + if name == id_column: + continue + entry, array, column_payloads = _encode_annotation_column( + annotations.column(name), name, num_rows, evidence_dict + ) + physical = f"{name}__end" if entry["kind"] == "multi" else name + if physical != name and physical in existing: + raise ValueError( + f"column '{name}' is multi-valued but '{physical}' already exists in " + "the annotations table; rename one of them" + ) + columns[physical] = array + manifest_columns[name] = entry + payloads.extend(column_payloads) + + if evidence_dict: + payloads.extend(_dict_payloads(_EVIDENCE_DICT_NAME, list(evidence_dict))) + + projections_table, projection_manifest = _encode_projections( + projections_metadata, projections_data, ids + ) + + manifest = { + "idColumn": id_column, + "columns": manifest_columns, + "projections": projection_manifest, + } + metadata = { + **{ + k: v + for k, v in (annotations.schema.metadata or {}).items() + if k != MANIFEST_KEY + }, + FORMAT_VERSION_KEY: str(CONTAINER_VERSION).encode(), + MANIFEST_KEY: json.dumps(manifest, separators=(",", ":")).encode(), + } + + payload_table = _required_table( + { + "name": pa.array([n for n, _ in payloads], type=pa.string()), + "data": pa.array([d for _, d in payloads], type=pa.binary()), + } + ) + + return ( + _write(_required_table(columns, metadata)), + _write(projections_metadata), + _write(projections_table), + _write(payload_table), + ) diff --git a/apps/protspace/tests/test_bundle_v3_encode.py b/apps/protspace/tests/test_bundle_v3_encode.py new file mode 100644 index 00000000..39f8c151 --- /dev/null +++ b/apps/protspace/tests/test_bundle_v3_encode.py @@ -0,0 +1,416 @@ +"""Encoder half of parquetbundle format v3 (``data/io/bundle_v3.encode_v3``). + +The assertions here pin the two contracts the encoder has to honour: + +* the *physical* contract that makes the browser's zero-copy read possible + (non-nullable, PLAIN, one row group, little-endian payload buffers), and +* the *semantic* contract that a v3 bundle must classify and order categories + exactly like the browser's v2 reader (``conversion.ts``), so both paths + produce the same colours and legend. +""" + +import io +import json + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from protspace.data.annotations.encoding import FORMAT_VERSION_KEY, stamp_format_version +from protspace.data.io.bundle_v3 import MANIFEST_KEY, encode_v3 + + +def make_annotations(**columns: list) -> pa.Table: + """Annotations table shaped like ``BaseProcessor._create_protein_annotations_table``.""" + n = len(next(iter(columns.values()))) + data = {"protein_id": [f"p{i}" for i in range(n)], **columns} + return stamp_format_version(pa.table(data)) + + +def make_projections(names_dims=(("A", 2),), ids=None): + ids = ids or ["p0"] + meta = pa.table( + { + "projection_name": [n for n, _ in names_dims], + "dimensions": [d for _, d in names_dims], + "info_json": ["{}"] * len(names_dims), + "source": [""] * len(names_dims), + } + ) + rows = [] + for name, dim in names_dims: + for i, pid in enumerate(ids): + rows.append( + { + "projection_name": name, + "identifier": pid, + "x": float(i), + "y": float(-i), + "z": float(i * 2) if dim == 3 else None, + } + ) + frame = pd.DataFrame(rows).astype({"x": "float32", "y": "float32", "z": "float32"}) + return meta, pa.Table.from_pandas(frame) + + +def encode(annotations: pa.Table, names_dims=(("A", 2),)): + """Encode ``annotations`` with matching projections; return the four parts.""" + ids = annotations.column("protein_id").to_pylist() + meta, data = make_projections(names_dims, ids) + return encode_v3(annotations, meta, data) + + +def read(part: bytes) -> pa.Table: + return pq.read_table(io.BytesIO(part)) + + +def manifest_of(part1: bytes) -> dict: + return json.loads(read(part1).schema.metadata[MANIFEST_KEY]) + + +def payloads_of(part6: bytes) -> dict[str, bytes]: + table = read(part6) + return dict( + zip( + table.column("name").to_pylist(), + table.column("data").to_pylist(), + strict=True, + ) + ) + + +def labels_of(payloads: dict[str, bytes], column: str) -> list[str]: + blob = payloads[f"dict:{column}"] + ends = np.frombuffer(payloads[f"dict:{column}:end"], " Date: Sun, 6 Sep 2026 00:23:20 +0200 Subject: [PATCH 02/31] feat(utils): read CSR annotation storage in the accessor funnel --- packages/utils/src/types.ts | 43 ++++++ .../annotation-data-access.test.ts | 138 ++++++++++++++++++ .../visualization/annotation-data-access.ts | 70 ++++++++- .../visualization/plot-data-accessors.test.ts | 80 +++++++++- .../src/visualization/plot-data-accessors.ts | 43 +++++- 5 files changed, 369 insertions(+), 5 deletions(-) diff --git a/packages/utils/src/types.ts b/packages/utils/src/types.ts index 6a38ea7b..dc111d45 100644 --- a/packages/utils/src/types.ts +++ b/packages/utils/src/types.ts @@ -44,6 +44,7 @@ export interface Annotation { * index, or `-1` when the protein has no value for this column. * - `SparseMultiValueAnnotationData`: compact single-value base plus overrides for the uncommon * multi-valued rows. + * - `CsrAnnotationData`: flat compressed-sparse-row codes, as delivered by bundle format v3. * - `(readonly number[])[]`: densely multi-valued column. `data[proteinIdx]` is the * list of indices; an empty array means missing. */ @@ -54,9 +55,43 @@ export interface SparseMultiValueAnnotationData { readonly length: number; } +/** + * Compressed sparse row storage for a multi-valued column (bundle format v3). + * + * Row `i` owns `codes[end[i - 1] .. end[i])`, with `end[-1]` conceptually 0, so + * a row with no values is `end[i - 1] === end[i]`. `end` is non-decreasing and + * `end[length - 1] === codes.length`. + */ +export interface CsrAnnotationData { + readonly kind: 'csr'; + readonly end: Int32Array; + readonly codes: Int32Array; + readonly length: number; +} + +/** + * Per-hit scores for a CSR column, indexed by the same hit numbering as + * {@link CsrAnnotationData.codes}: hit `h` owns `values[hitEnd[h - 1] .. hitEnd[h])` + * (`hitEnd[-1]` conceptually 0). An empty range means the hit carries no score. + */ +export interface CsrScores { + readonly hitEnd: Int32Array; + readonly values: Float32Array; +} + +/** + * Per-hit evidence for a CSR column, one code per hit: `-1` means none, + * otherwise the evidence string is `dict[code]`. + */ +export interface CsrEvidence { + readonly codes: Int32Array; + readonly dict: readonly string[]; +} + export type AnnotationData = | Int32Array | SparseMultiValueAnnotationData + | CsrAnnotationData | readonly (readonly number[])[]; /** A value transferred from a reference protein by Embedding Annotation Transfer (EAT). */ @@ -153,6 +188,14 @@ export interface VisualizationData { annotation_predicted?: AnnotationPredictedData; annotation_scores?: Record; annotation_evidence?: Record; + /** + * v3 counterparts of the two records above, flat per hit instead of nested per + * protein. Deliberately separate optional fields rather than a union with the + * nested form: the existing `annotation_scores?.[key]?.[i]` indexers stay valid, + * and a v1/v2 load never populates these. At most one form is present per column. + */ + annotation_scores_csr?: Record; + annotation_evidence_csr?: Record; /** * Raw projection-statistics parquet part (bundle part 5) as read, carried * unparsed so an export re-emits it instead of dropping it. This is the diff --git a/packages/utils/src/visualization/annotation-data-access.test.ts b/packages/utils/src/visualization/annotation-data-access.test.ts index 0a6c5057..3d183fc0 100644 --- a/packages/utils/src/visualization/annotation-data-access.test.ts +++ b/packages/utils/src/visualization/annotation-data-access.test.ts @@ -1,12 +1,15 @@ import { describe, it, expect } from 'vitest'; import { + getCsrHitRange, getProteinAnnotationIndices, getProteinAnnotationCount, getFirstAnnotationIndex, + isCsrAnnotationData, isMultilabelAnnotationData, isMultilabelAnnotationDataCached, sliceAnnotationData, } from './annotation-data-access'; +import type { CsrAnnotationData } from '../types'; describe('annotation-data-access', () => { describe('Int32Array storage', () => { @@ -134,6 +137,141 @@ describe('annotation-data-access', () => { }); }); +describe('CSR storage', () => { + // rows: 0 -> [5, 6], 1 -> [], 2 -> [2], 3 -> [], 4 -> [0, 1, 9] + // First and last rows carry hits, and `end[-1]` is conceptually 0, so row 0's + // range is [0, end[0]). + const csr = (): CsrAnnotationData => ({ + kind: 'csr', + end: Int32Array.from([2, 2, 3, 3, 6]), + codes: Int32Array.from([5, 6, 2, 0, 1, 9]), + length: 5, + }); + + it('is recognised without being mistaken for the other tagged shape', () => { + expect(isCsrAnnotationData(csr())).toBe(true); + expect(isCsrAnnotationData(Int32Array.from([0, 1]))).toBe(false); + expect(isCsrAnnotationData([[0], [1]])).toBe(false); + expect( + isCsrAnnotationData({ + kind: 'sparse-multi', + base: Int32Array.from([0]), + overrides: new Map(), + length: 1, + }), + ).toBe(false); + }); + + it('derives half-open hit ranges, empty outside the row count', () => { + const data = csr(); + expect(getCsrHitRange(data, 0)).toEqual([0, 2]); + expect(getCsrHitRange(data, 1)).toEqual([2, 2]); + expect(getCsrHitRange(data, 4)).toEqual([3, 6]); + expect(getCsrHitRange(data, 5)).toEqual([0, 0]); + expect(getCsrHitRange(data, -1)).toEqual([0, 0]); + }); + + it('returns indices for empty, single-hit, first and last rows', () => { + const data = csr(); + expect(getProteinAnnotationIndices(data, 0)).toEqual([5, 6]); + expect(getProteinAnnotationIndices(data, 1)).toEqual([]); + expect(getProteinAnnotationIndices(data, 2)).toEqual([2]); + expect(getProteinAnnotationIndices(data, 3)).toEqual([]); + expect(getProteinAnnotationIndices(data, 4)).toEqual([0, 1, 9]); + }); + + it('treats an empty first row as the [0, 0) range', () => { + const data: CsrAnnotationData = { + kind: 'csr', + end: Int32Array.from([0, 1]), + codes: Int32Array.from([3]), + length: 2, + }; + expect(getProteinAnnotationIndices(data, 0)).toEqual([]); + expect(getProteinAnnotationCount(data, 0)).toBe(0); + expect(getFirstAnnotationIndex(data, 0)).toBe(-1); + expect(getProteinAnnotationIndices(data, 1)).toEqual([3]); + }); + + it('returns a real Array, not a typed-array view', () => { + // Callers run `.map`/`.flatMap`/`.some` on the result; a subarray would only + // fail on `.flatMap`, so the shape itself is asserted. + const indices = getProteinAnnotationIndices(csr(), 0); + expect(Array.isArray(indices)).toBe(true); + expect(indices.flatMap((i) => [i, i])).toEqual([5, 5, 6, 6]); + }); + + it('counts hits per row', () => { + const data = csr(); + expect(getProteinAnnotationCount(data, 0)).toBe(2); + expect(getProteinAnnotationCount(data, 1)).toBe(0); + expect(getProteinAnnotationCount(data, 2)).toBe(1); + expect(getProteinAnnotationCount(data, 4)).toBe(3); + }); + + it('returns the first hit or -1', () => { + const data = csr(); + expect(getFirstAnnotationIndex(data, 0)).toBe(5); + expect(getFirstAnnotationIndex(data, 1)).toBe(-1); + expect(getFirstAnnotationIndex(data, 2)).toBe(2); + expect(getFirstAnnotationIndex(data, 4)).toBe(0); + }); + + it('matches the other shapes on out-of-range and negative indices', () => { + const data = csr(); + for (const idx of [5, 99, -1]) { + expect(getProteinAnnotationIndices(data, idx)).toEqual([]); + expect(getProteinAnnotationCount(data, idx)).toBe(0); + expect(getFirstAnnotationIndex(data, idx)).toBe(-1); + } + }); + + it('detects multilabel rows from the end deltas', () => { + expect(isMultilabelAnnotationData(csr())).toBe(true); + expect(isMultilabelAnnotationDataCached(csr())).toBe(true); + const singles: CsrAnnotationData = { + kind: 'csr', + end: Int32Array.from([1, 1, 2]), + codes: Int32Array.from([4, 7]), + length: 3, + }; + expect(isMultilabelAnnotationData(singles)).toBe(false); + expect(isMultilabelAnnotationDataCached(singles)).toBe(false); + }); + + it('slices to CSR, preserving hit order and dropping out-of-range rows', () => { + const sliced = sliceAnnotationData(csr(), [4, 1, 0, 99]); + expect(isCsrAnnotationData(sliced)).toBe(true); + const out = sliced as CsrAnnotationData; + expect(out.length).toBe(4); + expect(Array.from(out.end)).toEqual([3, 3, 5, 5]); + expect(Array.from(out.codes)).toEqual([0, 1, 9, 5, 6]); + expect(getProteinAnnotationIndices(out, 0)).toEqual([0, 1, 9]); + expect(getProteinAnnotationIndices(out, 1)).toEqual([]); + expect(getProteinAnnotationIndices(out, 2)).toEqual([5, 6]); + expect(getProteinAnnotationIndices(out, 3)).toEqual([]); + }); + + it('slices into fresh buffers that do not alias the source', () => { + const data = csr(); + const out = sliceAnnotationData(data, [0, 1, 2, 3, 4]) as CsrAnnotationData; + expect(Array.from(out.codes)).toEqual(Array.from(data.codes)); + expect(Array.from(out.end)).toEqual(Array.from(data.end)); + expect(out.codes.buffer).not.toBe(data.codes.buffer); + expect(out.end.buffer).not.toBe(data.end.buffer); + out.codes[0] = 42; + out.end[0] = 0; + expect(data.codes[0]).toBe(5); + expect(data.end[0]).toBe(2); + }); + + it('slices an all-empty selection to zero-length codes', () => { + const out = sliceAnnotationData(csr(), [1, 3]) as CsrAnnotationData; + expect(out.codes.length).toBe(0); + expect(Array.from(out.end)).toEqual([0, 0]); + }); +}); + describe('isMultilabelAnnotationDataCached', () => { it('agrees with the uncached form on every storage shape', () => { const dense: number[][] = [[0], [1, 2], [0]]; diff --git a/packages/utils/src/visualization/annotation-data-access.ts b/packages/utils/src/visualization/annotation-data-access.ts index 9d506d8d..39c67852 100644 --- a/packages/utils/src/visualization/annotation-data-access.ts +++ b/packages/utils/src/visualization/annotation-data-access.ts @@ -1,4 +1,8 @@ -import type { AnnotationData, SparseMultiValueAnnotationData } from '../types.js'; +import type { + AnnotationData, + CsrAnnotationData, + SparseMultiValueAnnotationData, +} from '../types.js'; export function isSparseMultiValueAnnotationData( data: AnnotationData, @@ -6,6 +10,29 @@ export function isSparseMultiValueAnnotationData( return 'kind' in data && data.kind === 'sparse-multi'; } +/** + * Tagged storage is checked before any `instanceof`: a CSR container is a plain + * object holding typed arrays, not a typed array itself. + */ +export function isCsrAnnotationData(data: AnnotationData): data is CsrAnnotationData { + return 'kind' in data && data.kind === 'csr'; +} + +/** + * Half-open `[start, stop)` range of hits owned by a protein in CSR storage. + * Out-of-range and negative protein indices yield an empty range. + * + * Also the hit numbering for the parallel `annotation_scores_csr` / + * `annotation_evidence_csr` payloads, which is why this is exported. + */ +export function getCsrHitRange( + data: CsrAnnotationData, + proteinIdx: number, +): readonly [number, number] { + if (proteinIdx < 0 || proteinIdx >= data.length) return [0, 0]; + return [proteinIdx === 0 ? 0 : data.end[proteinIdx - 1], data.end[proteinIdx]]; +} + /** * Memo for {@link isMultilabelAnnotationDataCached}. * @@ -43,6 +70,12 @@ export function isMultilabelAnnotationData(data: AnnotationData): boolean { } return false; } + if (isCsrAnnotationData(data)) { + for (let i = 0; i < data.length; i++) { + if (data.end[i] - (i === 0 ? 0 : data.end[i - 1]) > 1) return true; + } + return false; + } if (data instanceof Int32Array) return false; return data.some((values) => values.length > 1); } @@ -66,6 +99,12 @@ export function getProteinAnnotationIndices( const value = data.base[proteinIdx]; return value < 0 ? [] : [value]; } + if (isCsrAnnotationData(data)) { + // A subarray view would be cheaper, but callers run `.map`/`.flatMap`/`.some` + // on the result, so the contract stays `readonly number[]`. + const [start, stop] = getCsrHitRange(data, proteinIdx); + return start === stop ? [] : Array.from(data.codes.subarray(start, stop)); + } if (data instanceof Int32Array) { if (proteinIdx < 0 || proteinIdx >= data.length) return []; const value = data[proteinIdx]; @@ -82,6 +121,12 @@ export function getProteinAnnotationCount(data: AnnotationData, proteinIdx: numb if (proteinIdx < 0 || proteinIdx >= data.base.length) return 0; return data.base[proteinIdx] < 0 ? 0 : 1; } + if (isCsrAnnotationData(data)) { + // Range inlined rather than via getCsrHitRange: this and getFirstAnnotationIndex + // run per point per frame, and the tuple would be an allocation each. + if (proteinIdx < 0 || proteinIdx >= data.length) return 0; + return data.end[proteinIdx] - (proteinIdx === 0 ? 0 : data.end[proteinIdx - 1]); + } if (data instanceof Int32Array) { if (proteinIdx < 0 || proteinIdx >= data.length) return 0; return data[proteinIdx] < 0 ? 0 : 1; @@ -101,6 +146,11 @@ export function getFirstAnnotationIndex(data: AnnotationData, proteinIdx: number if (proteinIdx < 0 || proteinIdx >= data.base.length) return -1; return data.base[proteinIdx]; } + if (isCsrAnnotationData(data)) { + if (proteinIdx < 0 || proteinIdx >= data.length) return -1; + const start = proteinIdx === 0 ? 0 : data.end[proteinIdx - 1]; + return start === data.end[proteinIdx] ? -1 : data.codes[start]; + } if (data instanceof Int32Array) { if (proteinIdx < 0 || proteinIdx >= data.length) return -1; return data[proteinIdx]; @@ -126,6 +176,24 @@ export function sliceAnnotationData(data: AnnotationData, indices: number[]): An ? { kind: 'sparse-multi', base, overrides, length: base.length } : base; } + if (isCsrAnnotationData(data)) { + const end = new Int32Array(indices.length); + let total = 0; + for (let i = 0; i < indices.length; i++) { + const [start, stop] = getCsrHitRange(data, indices[i]); + total += stop - start; + end[i] = total; + } + const codes = new Int32Array(total); + let cursor = 0; + for (let i = 0; i < indices.length; i++) { + const [start, stop] = getCsrHitRange(data, indices[i]); + if (start === stop) continue; + codes.set(data.codes.subarray(start, stop), cursor); + cursor += stop - start; + } + return { kind: 'csr', end, codes, length: indices.length }; + } if (data instanceof Int32Array) { const out = new Int32Array(indices.length); for (let i = 0; i < indices.length; i++) { diff --git a/packages/utils/src/visualization/plot-data-accessors.test.ts b/packages/utils/src/visualization/plot-data-accessors.test.ts index ffbfef87..ba6b4649 100644 --- a/packages/utils/src/visualization/plot-data-accessors.test.ts +++ b/packages/utils/src/visualization/plot-data-accessors.test.ts @@ -8,7 +8,7 @@ import { getProteinEvidence, buildTooltipView, } from './plot-data-accessors'; -import type { VisualizationData } from '../types'; +import type { CsrAnnotationData, CsrEvidence, CsrScores, VisualizationData } from '../types'; const baseData = (): VisualizationData => ({ protein_ids: ['p0', 'p1', 'p2'], @@ -415,3 +415,81 @@ describe('plot-data-accessors', () => { }); }); }); + +describe('CSR score and evidence payloads', () => { + // species rows: p0 -> hits 0,1 ; p1 -> no hits ; p2 -> hit 2 + const csrRows: CsrAnnotationData = { + kind: 'csr', + end: Int32Array.from([2, 2, 3]), + codes: Int32Array.from([0, 1, 2]), + length: 3, + }; + // hit 0 -> [1.5]; hit 1 -> no scores; hit 2 -> [0.25, 0.5] + const csrScores: CsrScores = { + hitEnd: Int32Array.from([1, 1, 3]), + values: Float32Array.from([1.5, 0.25, 0.5]), + }; + const csrEvidence: CsrEvidence = { + codes: Int32Array.from([0, -1, 1]), + dict: ['IDA', 'ECO:1'], + }; + + const csrData = (): VisualizationData => { + const data = baseData(); + data.annotation_data.species = csrRows; + data.annotation_scores_csr = { species: csrScores }; + data.annotation_evidence_csr = { species: csrEvidence }; + return data; + }; + + it('reads scores per hit, null for a hit with no score values', () => { + const data = csrData(); + expect(getProteinScores(data, 0, 'species')).toEqual([[1.5], null]); + expect(getProteinScores(data, 1, 'species')).toEqual([]); + expect(getProteinScores(data, 2, 'species')).toEqual([[0.25, 0.5]]); + }); + + it('reads evidence per hit, null for code -1', () => { + const data = csrData(); + expect(getProteinEvidence(data, 0, 'species')).toEqual(['IDA', null]); + expect(getProteinEvidence(data, 1, 'species')).toEqual([]); + expect(getProteinEvidence(data, 2, 'species')).toEqual(['ECO:1']); + }); + + it('returns empty for out-of-range and negative protein indices', () => { + const data = csrData(); + for (const idx of [3, 99, -1]) { + expect(getProteinScores(data, idx, 'species')).toEqual([]); + expect(getProteinEvidence(data, idx, 'species')).toEqual([]); + } + }); + + it('prefers the nested records when both forms are present', () => { + const data = csrData(); + data.annotation_scores = { species: [[[9]], [], []] }; + data.annotation_evidence = { species: [['NESTED'], [], []] }; + expect(getProteinScores(data, 0, 'species')).toEqual([[9]]); + expect(getProteinEvidence(data, 0, 'species')).toEqual(['NESTED']); + }); + + it('ignores CSR payloads when the column storage is not CSR', () => { + // The flat payloads are numbered by CSR hit, so without CSR storage there is + // no hit range to index them by. + const data = csrData(); + data.annotation_data.species = Int32Array.of(0, 1, 2); + expect(getProteinScores(data, 0, 'species')).toEqual([]); + expect(getProteinEvidence(data, 0, 'species')).toEqual([]); + }); + + it('resolves annotation values through CSR storage', () => { + expect(getProteinAnnotationValues(csrData(), 0, 'species')).toEqual(['human', 'mouse']); + expect(getProteinAnnotationValues(csrData(), 1, 'species')).toEqual([]); + }); + + it('feeds the tooltip view without changing its shape', () => { + const view = buildTooltipView(csrData(), 0, 'species'); + expect(view.blocks[0].scores).toEqual([[1.5], null]); + expect(view.blocks[0].evidence).toEqual(['IDA', null]); + expect(view.blocks[0].displayValues).toEqual(['human', 'mouse']); + }); +}); diff --git a/packages/utils/src/visualization/plot-data-accessors.ts b/packages/utils/src/visualization/plot-data-accessors.ts index 4dd49070..00fbe3b6 100644 --- a/packages/utils/src/visualization/plot-data-accessors.ts +++ b/packages/utils/src/visualization/plot-data-accessors.ts @@ -1,5 +1,9 @@ import type { VisualizationData, NumericAnnotationType, PredictedCell } from '../types.js'; -import { getProteinAnnotationIndices } from './annotation-data-access.js'; +import { + getCsrHitRange, + getProteinAnnotationIndices, + isCsrAnnotationData, +} from './annotation-data-access.js'; import { isAutoClusterColumnName } from './annotation-statistics.js'; import { getPredictedCell, getPredictedCellValues } from './eat-overlay.js'; import { getNumericBinLabelMap } from './numeric-binning.js'; @@ -59,13 +63,37 @@ export function getProteinNumericType( return annotation?.numericType ?? annotation?.numericMetadata?.numericType ?? 'float'; } +/** + * Hit range this protein owns, for the flat v3 score/evidence payloads. Empty + * unless the column's storage is CSR — the flat payloads are numbered by CSR hit, + * so there is nothing to index them by otherwise. + */ +function getCsrHitRangeFor( + data: VisualizationData, + proteinIdx: number, + annotationKey: string, +): readonly [number, number] { + const rows = data.annotation_data?.[annotationKey]; + return rows && isCsrAnnotationData(rows) ? getCsrHitRange(rows, proteinIdx) : [0, 0]; +} + export function getProteinScores( data: VisualizationData, proteinIdx: number, annotationKey: string, ): (number[] | null)[] { const scores = data.annotation_scores?.[annotationKey]?.[proteinIdx]; - return Array.isArray(scores) ? scores : []; + if (Array.isArray(scores)) return scores; + const csr = data.annotation_scores_csr?.[annotationKey]; + if (!csr) return []; + const [start, stop] = getCsrHitRangeFor(data, proteinIdx, annotationKey); + const out: (number[] | null)[] = []; + for (let hit = start; hit < stop; hit++) { + const from = hit === 0 ? 0 : csr.hitEnd[hit - 1]; + const to = csr.hitEnd[hit]; + out.push(to > from ? Array.from(csr.values.subarray(from, to)) : null); + } + return out; } export function getProteinEvidence( @@ -74,7 +102,16 @@ export function getProteinEvidence( annotationKey: string, ): (string | null)[] { const evidence = data.annotation_evidence?.[annotationKey]?.[proteinIdx]; - return Array.isArray(evidence) ? evidence : []; + if (Array.isArray(evidence)) return evidence; + const csr = data.annotation_evidence_csr?.[annotationKey]; + if (!csr) return []; + const [start, stop] = getCsrHitRangeFor(data, proteinIdx, annotationKey); + const out: (string | null)[] = []; + for (let hit = start; hit < stop; hit++) { + const code = csr.codes[hit]; + out.push(code >= 0 ? (csr.dict[code] ?? null) : null); + } + return out; } /** From 4e4d9ee1a26e2ad9d8d19c28d6ae73e99e2a0396 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 00:24:17 +0200 Subject: [PATCH 03/31] feat(utils): handle CSR annotation storage at the bypass sites --- .../utils/bundle-roundtrip.test.ts | 67 ++++++++++ .../data-loader/utils/conversion.test.ts | 122 ++++++++++++++++++ .../data-loader/utils/conversion.ts | 38 +++++- .../styling/visibility-model.test.ts | 35 +++++ .../scatter-plot/styling/visibility-model.ts | 26 +++- packages/utils/src/parquet/bundle-writer.ts | 12 +- .../src/visualization/eat-overlay.test.ts | 77 +++++++++++ .../utils/src/visualization/eat-overlay.ts | 37 ++++++ .../slice-visualization-data.test.ts | 68 ++++++++++ .../visualization/slice-visualization-data.ts | 69 +++++++++- 10 files changed, 543 insertions(+), 8 deletions(-) diff --git a/packages/core/src/components/data-loader/utils/bundle-roundtrip.test.ts b/packages/core/src/components/data-loader/utils/bundle-roundtrip.test.ts index 479b85e3..7658497a 100644 --- a/packages/core/src/components/data-loader/utils/bundle-roundtrip.test.ts +++ b/packages/core/src/components/data-loader/utils/bundle-roundtrip.test.ts @@ -9,6 +9,9 @@ import { import { createParquetBundle, countBundleDelimiters, + getProteinAnnotationValues, + getProteinEvidence, + getProteinScores, findBundleDelimiterPositions, isParquetBundle, type Annotation, @@ -712,3 +715,67 @@ describe('numeric annotation type fidelity', () => { expect(reimported.annotation_data.length).toBeUndefined(); }); }); + +describe('export from CSR storage (bundle format v3 in memory)', () => { + /** + * The writer only ever stamps format version 2, so a CSR-loaded dataset has to + * serialize to exactly the same v2 cells as the equivalent nested dataset. + * p0 → 'a|IDA', p1 → nothing, p2 → 'b|0.5,0.25;a'. + */ + const annotations: Record = { + fam: { + kind: 'categorical', + values: ['a', 'b'], + colors: ['#000', '#fff'], + shapes: ['circle', 'square'], + }, + }; + const shared = { + protein_ids: ['p0', 'p1', 'p2'], + projections: [{ name: 'umap', dimension: 2 as const, data: new Float32Array(6) }], + annotations, + }; + + const nested: VisualizationData = { + ...shared, + annotation_data: { fam: [[0], [], [1, 0]] }, + annotation_scores: { fam: [[null], [], [[0.5, 0.25], null]] }, + annotation_evidence: { fam: [['IDA'], [], [null, null]] }, + }; + const csr: VisualizationData = { + ...shared, + annotation_data: { + fam: { kind: 'csr', end: Int32Array.of(1, 1, 3), codes: Int32Array.of(0, 1, 0), length: 3 }, + }, + annotation_scores_csr: { + fam: { hitEnd: Int32Array.of(0, 2, 2), values: Float32Array.of(0.5, 0.25) }, + }, + annotation_evidence_csr: { + fam: { codes: Int32Array.of(0, -1, -1), dict: ['IDA'] }, + }, + }; + + it('writes the same annotation cells as the equivalent nested dataset', async () => { + const readCells = async (data: VisualizationData) => { + const extraction = await extractRowsFromParquetBundle(createParquetBundle(data)); + return convertParquetToVisualizationData(extraction); + }; + const fromCsr = await readCells(csr); + const fromNested = await readCells(nested); + + expect(fromCsr.protein_ids).toEqual(fromNested.protein_ids); + expect(fromCsr.annotations.fam.values).toEqual(fromNested.annotations.fam.values); + for (let i = 0; i < 3; i++) { + expect(getProteinAnnotationValues(fromCsr, i, 'fam')).toEqual( + getProteinAnnotationValues(fromNested, i, 'fam'), + ); + expect(getProteinScores(fromCsr, i, 'fam')).toEqual(getProteinScores(fromNested, i, 'fam')); + expect(getProteinEvidence(fromCsr, i, 'fam')).toEqual( + getProteinEvidence(fromNested, i, 'fam'), + ); + } + // Sanity: the fixture actually carries scores and evidence through the export. + expect(getProteinEvidence(fromCsr, 0, 'fam')).toEqual(['IDA']); + expect(getProteinScores(fromCsr, 2, 'fam')).toEqual([[0.5, 0.25], null]); + }); +}); diff --git a/packages/core/src/components/data-loader/utils/conversion.test.ts b/packages/core/src/components/data-loader/utils/conversion.test.ts index 13dc3a0d..f5a1a387 100644 --- a/packages/core/src/components/data-loader/utils/conversion.test.ts +++ b/packages/core/src/components/data-loader/utils/conversion.test.ts @@ -3,12 +3,15 @@ import { readFileSync } from 'node:fs'; import { createParquetBundle, getProteinAnnotationIndices, + isCsrAnnotationData, materializeEatOverlay, materializeVisualizationData, } from '@protspace/utils'; +import type { CsrAnnotationData, VisualizationData } from '@protspace/utils'; import { convertParquetToVisualizationData, convertParquetToVisualizationDataOptimized, + normalizeEatCompanionColumns, parseAnnotationValue, splitCategoricalAnnotationValues, } from './conversion'; @@ -735,3 +738,122 @@ describe('splitCategoricalAnnotationValues v2', () => { expect(splitCategoricalAnnotationValues(raw, 2)).toEqual(['A (n%3B1)|1', 'B (n%3B2)|2']); }); }); + +describe('normalizeEatCompanionColumns over CSR storage (bundle format v3)', () => { + /** + * Base column `ec` in CSR, values `['A', null, 'B']` — the `null` slot is what makes + * the remap drop a code, so the rebuilt payload has to compact: + * p0 → [] (curated missing → the EAT companion applies) + * p1 → [A] + * p2 → [null, B] + * p3 → [A, B] + * The companion `ec__pred_value` is CSR too, with its scores/evidence in the flat + * v3 payloads rather than the nested records. + */ + function csrEatData(): VisualizationData { + return { + protein_ids: ['p0', 'p1', 'p2', 'p3'], + projections: [{ name: 'umap', dimension: 2, data: new Float32Array(8) }], + annotations: { + ec: { + kind: 'categorical', + values: ['A', null, 'B'], + colors: ['#f00', '#0f0', '#00f'], + shapes: ['circle', 'circle', 'circle'], + }, + ec__pred_value: { + kind: 'categorical', + values: ['C'], + colors: ['#fff'], + shapes: ['circle'], + }, + ec__pred_confidence: { kind: 'numeric', numericType: 'float', values: [] }, + ec__pred_source: { + kind: 'categorical', + values: ['p1'], + colors: ['#fff'], + shapes: ['circle'], + }, + }, + annotation_data: { + ec: { + kind: 'csr', + end: Int32Array.of(0, 1, 3, 5), + codes: Int32Array.of(0, 1, 2, 0, 2), + length: 4, + }, + ec__pred_value: { + kind: 'csr', + end: Int32Array.of(1, 1, 1, 1), + codes: Int32Array.of(0), + length: 4, + }, + ec__pred_source: { + kind: 'csr', + end: Int32Array.of(1, 1, 1, 1), + codes: Int32Array.of(0), + length: 4, + }, + }, + numeric_annotation_data: { ec__pred_confidence: [0.9, null, null, null] }, + annotation_scores_csr: { + ec__pred_value: { hitEnd: Int32Array.of(2), values: Float32Array.of(0.5, 0.25) }, + }, + annotation_evidence_csr: { + ec__pred_value: { codes: Int32Array.of(0), dict: ['IDA'] }, + }, + }; + } + + it('remaps CSR storage into fresh compacted buffers', () => { + const src = csrEatData(); + const sourceRows = src.annotation_data.ec as CsrAnnotationData; + const out = normalizeEatCompanionColumns(src); + const rows = out.annotation_data.ec; + if (!isCsrAnnotationData(rows)) throw new Error('expected CSR storage'); + + // 'A','B' survive, the null slot is dropped, 'C' is appended by the transfer. + expect(out.annotations.ec.values).toEqual(['A', 'B', 'C']); + // The transfer itself stays in `annotation_predicted`; only the curated rows + // are remapped here (materializeEatOverlay applies the prediction later). + expect(getProteinAnnotationIndices(rows, 0)).toEqual([]); + expect(getProteinAnnotationIndices(rows, 1)).toEqual([0]); // 'A' + expect(getProteinAnnotationIndices(rows, 2)).toEqual([1]); // null dropped, 'B' kept + expect(getProteinAnnotationIndices(rows, 3)).toEqual([0, 1]); + expect(Array.from(rows.end)).toEqual([0, 1, 2, 4]); + + // Compacted: one code fewer than the source, and the trailing slack is not + // retained — `.slice(0, written)` hands back an exactly-sized fresh buffer. + expect(rows.codes.length).toBe(4); + expect(rows.codes.buffer.byteLength).toBe(4 * 4); + expect(rows.codes.buffer).not.toBe(sourceRows.codes.buffer); + expect(Array.from(sourceRows.codes)).toEqual([0, 1, 2, 0, 2]); + }); + + it('reads the companion column scores/evidence from the flat v3 payloads', () => { + const cell = normalizeEatCompanionColumns(csrEatData()).annotation_predicted?.ec?.[0]; + expect(cell).toMatchObject({ + value: 'C', + confidence: 0.9, + source: 'p1', + scores: [[0.5, 0.25]], + evidence: ['IDA'], + }); + }); + + it('drops the companion columns from the flat payload records too', () => { + const out = normalizeEatCompanionColumns(csrEatData()); + expect(out.annotation_scores_csr).toEqual({}); + expect(out.annotation_evidence_csr).toEqual({}); + expect(out.annotation_data.ec__pred_value).toBeUndefined(); + }); + + it('does not grow the flat payload records on a dataset that has none', () => { + const src = csrEatData(); + delete src.annotation_scores_csr; + delete src.annotation_evidence_csr; + const out = normalizeEatCompanionColumns(src); + expect('annotation_scores_csr' in out).toBe(false); + expect('annotation_evidence_csr' in out).toBe(false); + }); +}); diff --git a/packages/core/src/components/data-loader/utils/conversion.ts b/packages/core/src/components/data-loader/utils/conversion.ts index 0ddcf0a3..f63cbd22 100644 --- a/packages/core/src/components/data-loader/utils/conversion.ts +++ b/packages/core/src/components/data-loader/utils/conversion.ts @@ -9,6 +9,9 @@ import { COLOR_SCHEMES, getEatConfidenceAnnotationKey, getProteinAnnotationIndices, + getProteinEvidence, + getProteinScores, + isCsrAnnotationData, isSparseMultiValueAnnotationData, isCuratedAnnotationMissing, isNAValue, @@ -222,6 +225,23 @@ function remapCategoricalStorage( } return { kind: 'sparse-multi', base, overrides, length: base.length }; } + if (isCsrAnnotationData(source)) { + // Remapping can drop a hit (`remap` → -1), so the row boundaries shift and the + // payload shrinks: both arrays are rebuilt from scratch, and `codes` is trimmed + // to what was written so the worst-case buffer is neither retained nor shared. + const end = new Int32Array(source.length); + const codes = new Int32Array(source.codes.length); + let written = 0; + for (let i = 0; i < source.length; i++) { + const stop = source.end[i]; + for (let hit = i === 0 ? 0 : source.end[i - 1]; hit < stop; hit++) { + const index = remap(source.codes[hit]); + if (index >= 0) codes[written++] = index; + } + end[i] = written; + } + return { kind: 'csr', end, codes: codes.slice(0, written), length: source.length }; + } return source.map((indices) => indices.map(remap).filter((index) => index >= 0)); } @@ -230,7 +250,7 @@ function remapCategoricalStorage( * This is intentionally the final conversion-boundary step so small, optimized, and separated * decoder paths all share exactly one validity rule. */ -function normalizeEatCompanionColumns(data: VisualizationData): VisualizationData { +export function normalizeEatCompanionColumns(data: VisualizationData): VisualizationData { const groups = new Map>>(); const reservedColumns = new Set(); for (const column of Object.keys(data.annotations)) { @@ -248,6 +268,10 @@ function normalizeEatCompanionColumns(data: VisualizationData): VisualizationDat const numeric_annotation_data = { ...data.numeric_annotation_data }; const annotation_scores = { ...data.annotation_scores }; const annotation_evidence = { ...data.annotation_evidence }; + // v3 columns carry their scores/evidence flat instead of nested; a companion + // column must lose both forms or the dropped column's payload outlives it. + const annotation_scores_csr = { ...data.annotation_scores_csr }; + const annotation_evidence_csr = { ...data.annotation_evidence_csr }; const annotation_predicted = { ...data.annotation_predicted }; for (const column of reservedColumns) { @@ -256,6 +280,8 @@ function normalizeEatCompanionColumns(data: VisualizationData): VisualizationDat delete numeric_annotation_data[column]; delete annotation_scores[column]; delete annotation_evidence[column]; + delete annotation_scores_csr[column]; + delete annotation_evidence_csr[column]; } for (const [base, group] of groups) { @@ -281,8 +307,9 @@ function normalizeEatCompanionColumns(data: VisualizationData): VisualizationDat if (!isCuratedAnnotationMissing(data, base, i)) continue; const values = readCategoricalStorageValues(data, group.value, i); const value = values.length > 0 ? values.join(';') : null; - const scores = data.annotation_scores?.[group.value]?.[i] ?? []; - const evidence = data.annotation_evidence?.[group.value]?.[i] ?? []; + // Via the accessors so a CSR-stored companion column resolves too. + const scores = getProteinScores(data, i, group.value); + const evidence = getProteinEvidence(data, i, group.value); const source = readCategoricalStorageValue(data, group.source, i); const confidence = confidences[i]; if ( @@ -374,6 +401,9 @@ function normalizeEatCompanionColumns(data: VisualizationData): VisualizationDat numeric_annotation_data, annotation_scores, annotation_evidence, + // Spread conditionally: a v1/v2 dataset never had these keys and must not grow them. + ...(data.annotation_scores_csr ? { annotation_scores_csr } : {}), + ...(data.annotation_evidence_csr ? { annotation_evidence_csr } : {}), annotation_predicted: Object.keys(annotation_predicted).length > 0 ? annotation_predicted : undefined, }; @@ -692,6 +722,8 @@ function restoreDeclaredNumericAnnotations( delete data.annotation_data[column]; delete data.annotation_scores?.[column]; delete data.annotation_evidence?.[column]; + delete data.annotation_scores_csr?.[column]; + delete data.annotation_evidence_csr?.[column]; } return data; } diff --git a/packages/core/src/components/scatter-plot/styling/visibility-model.test.ts b/packages/core/src/components/scatter-plot/styling/visibility-model.test.ts index 6213ec80..a1e92ff5 100644 --- a/packages/core/src/components/scatter-plot/styling/visibility-model.test.ts +++ b/packages/core/src/components/scatter-plot/styling/visibility-model.test.ts @@ -189,6 +189,41 @@ describe('computeVisibilityModel', () => { }); }); + // ── CSR storage (bundle format v3) — same rule 3/4 semantics ────────────── + describe('CSR storage: hidden iff every code hidden', () => { + // p0 → [A, B], p1 → [C], p2 → [] (empty row), p3 → [A] + const csr = (): AnnotationData => ({ + kind: 'csr', + end: Int32Array.of(2, 3, 3, 4), + codes: Int32Array.of(0, 1, 2, 0), + length: 4, + }); + + it('partially hidden CSR row stays visible', () => { + const data = makeData(['A', 'B', 'C'], csr()); + const model = computeVisibilityModel(baseInputs({ data, hiddenAnnotationValues: ['A'] })); + expect(model.opacityOf(point('p0', 0))).toBe(OPACITIES.base); + }); + + it('CSR row with every code hidden → opacity 0', () => { + // C stays visible so the all-hidden escape hatch cannot fire. + const data = makeData(['A', 'B', 'C'], csr()); + const model = computeVisibilityModel( + baseInputs({ data, hiddenAnnotationValues: ['A', 'B'] }), + ); + expect(model.opacityOf(point('p0', 0))).toBe(0); + expect(model.opacityOf(point('p3', 3))).toBe(0); + expect(model.opacityOf(point('p1', 1))).toBe(OPACITIES.base); + }); + + it('empty CSR row → hidden even with an empty hidden set', () => { + const data = makeData(['A', 'B', 'C'], csr()); + const model = computeVisibilityModel(baseInputs({ data, hiddenAnnotationValues: [] })); + expect(model.opacityOf(point('p2', 2))).toBe(0); + expect(model.opacityOf(point('p0', 0))).toBe(OPACITIES.base); + }); + }); + // ── Rule 5 — all-hidden escape hatch ────────────────────────────────────── describe('rule 5: all-hidden escape hatch rescues opacity (not colors)', () => { it('every value hidden → opacity falls back to base tier, allHidden true', () => { diff --git a/packages/core/src/components/scatter-plot/styling/visibility-model.ts b/packages/core/src/components/scatter-plot/styling/visibility-model.ts index d729ea0a..27e32b0b 100644 --- a/packages/core/src/components/scatter-plot/styling/visibility-model.ts +++ b/packages/core/src/components/scatter-plot/styling/visibility-model.ts @@ -40,7 +40,11 @@ import type { PlotDataPoint, VisualizationData, } from '@protspace/utils'; -import { isSparseMultiValueAnnotationData, toInternalValue } from '@protspace/utils'; +import { + isCsrAnnotationData, + isSparseMultiValueAnnotationData, + toInternalValue, +} from '@protspace/utils'; export interface VisibilityInputs { /** MATERIALIZED, un-query-filtered display data (keeps global indices). */ @@ -144,6 +148,26 @@ function buildHiddenMask( } mask[i] = everyHidden; } + } else if (isCsrAnnotationData(annotationRows)) { + const { end, codes, length: len } = annotationRows; + for (let i = 0; i < n; i++) { + if (i >= len) { + mask[i] = 1; + continue; + } + // Hit range inlined rather than via `getCsrHitRange`: the tuple would be an + // allocation per point, and this pass runs on every legend hide. + const stop = end[i]; + let everyHidden = 1; + // Empty row (start === stop) leaves this 1 — vacuously hidden, as `[].every()`. + for (let k = i === 0 ? 0 : end[i - 1]; k < stop; k++) { + if (isBinHidden(codes[k]) === 0) { + everyHidden = 0; + break; + } + } + mask[i] = everyHidden; + } } else { const len = annotationRows.length; for (let i = 0; i < n; i++) { diff --git a/packages/utils/src/parquet/bundle-writer.ts b/packages/utils/src/parquet/bundle-writer.ts index 48f2243d..796f005f 100644 --- a/packages/utils/src/parquet/bundle-writer.ts +++ b/packages/utils/src/parquet/bundle-writer.ts @@ -28,6 +28,7 @@ import { assertNoBundleDelimiter } from './delimiter-utils'; import { bigIntReplacer } from './bigint-utils'; import { isNumericAnnotation } from '../visualization/numeric-binning.js'; import { getProteinAnnotationIndices } from '../visualization/annotation-data-access.js'; +import { getProteinEvidence, getProteinScores } from '../visualization/plot-data-accessors.js'; import { getEatCompanionColumn, getPredictedCellValues } from '../visualization/eat-overlay.js'; import { isNAValue } from '../visualization/missing-values.js'; import { encodeAnnotationField } from './annotation-codec.js'; @@ -143,6 +144,13 @@ function createAnnotationsParquet(data: VisualizationData): ArrayBuffer { // Reconstruct the v2 wire cell positionally so decoded labels, evidence, and scores survive // export without structural semicolons/pipes being reinterpreted on reload. + // + // Read once per protein through the accessors rather than indexing the nested + // records: a v3 (CSR) dataset carries these flat per hit, and only the accessors + // know both layouts. Both return per-hit arrays in `getProteinAnnotationIndices` + // order, so `cellIndex` still lines up. + const proteinEvidence = getProteinEvidence(data, i, annotationName); + const proteinScores = getProteinScores(data, i, annotationName); const cellValues = getProteinAnnotationIndices(annotationIndices, i).flatMap( (valueIndex, cellIndex) => { const value = annotation.values[valueIndex]; @@ -152,8 +160,8 @@ function createAnnotationsParquet(data: VisualizationData): ArrayBuffer { // (it is not a MISSING_VALUE_TOKEN) and leak into downstream `protspace` tooling. // Mirrors the read side's readCategoricalStorageValues. if (value == null || isNAValue(value)) return []; - const evidence = data.annotation_evidence?.[annotationName]?.[i]?.[cellIndex]; - const scores = data.annotation_scores?.[annotationName]?.[i]?.[cellIndex]; + const evidence = proteinEvidence[cellIndex]; + const scores = proteinScores[cellIndex]; return [serializeCategoricalValue(value, evidence, scores)]; }, ); diff --git a/packages/utils/src/visualization/eat-overlay.test.ts b/packages/utils/src/visualization/eat-overlay.test.ts index 71f47503..22970238 100644 --- a/packages/utils/src/visualization/eat-overlay.test.ts +++ b/packages/utils/src/visualization/eat-overlay.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest'; import type { VisualizationData } from '../types'; import { getProteinAnnotationIndices, + isCsrAnnotationData, isSparseMultiValueAnnotationData, } from './annotation-data-access'; import { @@ -169,3 +170,79 @@ describe('clampReliabilityBound', () => { expect(clampReliabilityBound(0.42)).toBe(0.42); }); }); + +describe('EAT overlay over CSR storage (bundle format v3)', () => { + // p0 → [A], p1 → [] (curated missing), p2 → [B, C], p3 → [] (curated missing) + function createCsrData(): VisualizationData { + return { + protein_ids: ['p0', 'p1', 'p2', 'p3'], + projections: [{ name: 'umap', dimension: 2, data: new Float32Array(8) }], + annotations: { + ec: { + kind: 'categorical', + values: ['A', 'B', 'C'], + colors: ['#f00', '#0f0', '#00f'], + shapes: ['circle', 'circle', 'circle'], + }, + }, + annotation_data: { + ec: { + kind: 'csr', + end: Int32Array.of(1, 1, 3, 3), + codes: Int32Array.of(0, 1, 2), + length: 4, + }, + }, + annotation_predicted: { + // p1 gets a single-valued transfer, p3 a two-valued one. + ec: [ + null, + { value: 'C', confidence: 0.8, source: 'p0' }, + null, + { value: 'A;B', values: ['A', 'B'], confidence: 0.6, source: 'p2' }, + ], + }, + }; + } + + it('rebuilds CSR with predicted rows replaced and every other row preserved', () => { + const data = createCsrData(); + const out = materializeEatOverlay(data, 'ec', true); + const rows = out.annotation_data.ec; + expect(isCsrAnnotationData(rows)).toBe(true); + + expect(getProteinAnnotationIndices(rows, 0)).toEqual([0]); // untouched + expect(getProteinAnnotationIndices(rows, 1)).toEqual([2]); // predicted 'C' + expect(getProteinAnnotationIndices(rows, 2)).toEqual([1, 2]); // untouched + expect(getProteinAnnotationIndices(rows, 3)).toEqual([0, 1]); // predicted 'A;B' + + if (!isCsrAnnotationData(rows)) throw new Error('expected CSR'); + expect(Array.from(rows.end)).toEqual([1, 2, 4, 6]); + expect(rows.codes.length).toBe(6); + expect(rows.end[rows.length - 1]).toBe(rows.codes.length); + }); + + it('does not mutate or alias the curated CSR storage', () => { + const data = createCsrData(); + const source = data.annotation_data.ec; + const out = materializeEatOverlay(data, 'ec', true); + if (!isCsrAnnotationData(source) || !isCsrAnnotationData(out.annotation_data.ec)) { + throw new Error('expected CSR'); + } + expect(Array.from(source.end)).toEqual([1, 1, 3, 3]); + expect(Array.from(source.codes)).toEqual([0, 1, 2]); + expect(out.annotation_data.ec.codes.buffer).not.toBe(source.codes.buffer); + }); + + it('leaves the row alone when a prediction resolves to no known value', () => { + const data = createCsrData(); + data.annotation_predicted!.ec = [ + { value: 'not-a-category', confidence: 0.9, source: 'p2' }, + null, + null, + null, + ]; + const rows = materializeEatOverlay(data, 'ec', true).annotation_data.ec; + expect(getProteinAnnotationIndices(rows, 0)).toEqual([0]); + }); +}); diff --git a/packages/utils/src/visualization/eat-overlay.ts b/packages/utils/src/visualization/eat-overlay.ts index 10bad58d..9d5f338a 100644 --- a/packages/utils/src/visualization/eat-overlay.ts +++ b/packages/utils/src/visualization/eat-overlay.ts @@ -2,6 +2,7 @@ import type { AnnotationData, PredictedCell, VisualizationData } from '../types. import { getFirstAnnotationIndex, getProteinAnnotationIndices, + isCsrAnnotationData, isSparseMultiValueAnnotationData, } from './annotation-data-access.js'; import { isNAValue } from './missing-values.js'; @@ -246,6 +247,42 @@ function cloneWithPredictions( return { kind: 'sparse-multi', base, overrides, length: base.length }; } + if (isCsrAnnotationData(source)) { + // Rebuilt as CSR, never densified to `number[][]`: a 573K column would otherwise + // become one array per protein. Predicted rows replace the source row's hits + // wholesale; a prediction that resolves to nothing leaves the row alone, as the + // Int32Array and sparse-multi branches above do (CSR has no way to store the + // `-1` the dense branch falls back to). + const replacements = new Map(); + for (let i = 0; i < predictedCells.length && i < source.length; i++) { + const cell = predictedCells[i]; + if (!cell) continue; + const indices = predictedIndices(cell, valueToIndex); + if (indices.length > 0) replacements.set(i, indices); + } + const end = new Int32Array(source.length); + let total = 0; + for (let i = 0; i < source.length; i++) { + const replacement = replacements.get(i); + total += replacement ? replacement.length : source.end[i] - (i === 0 ? 0 : source.end[i - 1]); + end[i] = total; + } + const codes = new Int32Array(total); + let cursor = 0; + for (let i = 0; i < source.length; i++) { + const replacement = replacements.get(i); + if (replacement) { + for (const index of replacement) codes[cursor++] = index; + continue; + } + const start = i === 0 ? 0 : source.end[i - 1]; + const stop = source.end[i]; + if (stop > start) codes.set(source.codes.subarray(start, stop), cursor); + cursor += stop - start; + } + return { kind: 'csr', end, codes, length: source.length }; + } + const clone = source.slice(); for (let i = 0; i < predictedCells.length; i++) { const cell = predictedCells[i]; diff --git a/packages/utils/src/visualization/slice-visualization-data.test.ts b/packages/utils/src/visualization/slice-visualization-data.test.ts index 0a3b604c..c804aba8 100644 --- a/packages/utils/src/visualization/slice-visualization-data.test.ts +++ b/packages/utils/src/visualization/slice-visualization-data.test.ts @@ -1,5 +1,7 @@ import { describe, it, expect } from 'vitest'; import { sliceVisualizationDataByIndices } from './slice-visualization-data'; +import { getProteinAnnotationIndices, isCsrAnnotationData } from './annotation-data-access'; +import { getProteinEvidence, getProteinScores } from './plot-data-accessors'; import type { Annotation, VisualizationData } from '../types'; function baseViz(): VisualizationData { @@ -85,3 +87,69 @@ describe('sliceVisualizationDataByIndices', () => { expect(out.annotations).toBe(src.annotations); }); }); + +describe('sliceVisualizationDataByIndices over CSR storage (bundle format v3)', () => { + // p0 → [a], p1 → [] , p2 → [b, a], p3 → [b] + function csrViz(): VisualizationData { + return { + protein_ids: ['p0', 'p1', 'p2', 'p3'], + projections: [{ name: 'umap', dimension: 2, data: new Float32Array(8) }], + annotations: { + fam: { + kind: 'categorical', + values: ['a', 'b'], + colors: ['#000', '#fff'], + shapes: ['circle', 'square'], + }, + }, + annotation_data: { + fam: { + kind: 'csr', + end: Int32Array.of(1, 1, 3, 4), + codes: Int32Array.of(0, 1, 0, 1), + length: 4, + }, + }, + // Hits 0..3 in the same order as `codes`. + annotation_scores_csr: { + fam: { + hitEnd: Int32Array.of(1, 1, 3, 4), + values: Float32Array.of(0.5, 1.5, 2.5, 3.5), + }, + }, + annotation_evidence_csr: { + fam: { codes: Int32Array.of(0, -1, 1, 2), dict: ['IDA', 'IEA', 'IPI'] }, + }, + }; + } + + it('keeps the flat score/evidence payloads aligned with the sliced CSR rows', () => { + const src = csrViz(); + // Reversed order and a dropped row, so a slice that just copies would be wrong. + const out = sliceVisualizationDataByIndices(src, [3, 2, 1]); + + expect(isCsrAnnotationData(out.annotation_data.fam)).toBe(true); + expect(getProteinAnnotationIndices(out.annotation_data.fam, 0)).toEqual([1]); // was p3 + expect(getProteinAnnotationIndices(out.annotation_data.fam, 1)).toEqual([1, 0]); // was p2 + expect(getProteinAnnotationIndices(out.annotation_data.fam, 2)).toEqual([]); // was p1 + + // Same answers the source gave for the same proteins. + for (const [before, after] of [ + [3, 0], + [2, 1], + [1, 2], + ]) { + expect(getProteinScores(out, after, 'fam')).toEqual(getProteinScores(src, before, 'fam')); + expect(getProteinEvidence(out, after, 'fam')).toEqual(getProteinEvidence(src, before, 'fam')); + } + expect(getProteinScores(src, 2, 'fam')).toEqual([null, [1.5, 2.5]]); + expect(getProteinEvidence(src, 2, 'fam')).toEqual([null, 'IEA']); + expect(getProteinEvidence(src, 3, 'fam')).toEqual(['IPI']); + }); + + it('omits the flat payloads when the source has none', () => { + const out = sliceVisualizationDataByIndices(baseViz(), [0]); + expect(out.annotation_scores_csr).toBeUndefined(); + expect(out.annotation_evidence_csr).toBeUndefined(); + }); +}); diff --git a/packages/utils/src/visualization/slice-visualization-data.ts b/packages/utils/src/visualization/slice-visualization-data.ts index c256a062..a30a4f6a 100644 --- a/packages/utils/src/visualization/slice-visualization-data.ts +++ b/packages/utils/src/visualization/slice-visualization-data.ts @@ -1,5 +1,9 @@ -import type { VisualizationData } from '../types.js'; -import { sliceAnnotationData } from './annotation-data-access.js'; +import type { AnnotationData, CsrEvidence, CsrScores, VisualizationData } from '../types.js'; +import { + getCsrHitRange, + isCsrAnnotationData, + sliceAnnotationData, +} from './annotation-data-access.js'; /** * Build a VisualizationData constrained to `keptIndices` (ascending positions into @@ -28,6 +32,65 @@ export function sliceVisualizationDataByIndices( ? Object.fromEntries(Object.entries(src).map(([name, rows]) => [name, sliceRows(rows)])) : undefined; + /** + * Source hit numbers the kept proteins own, in kept order — the same order + * `sliceAnnotationData` concatenates their codes in, so the sliced flat payloads + * stay aligned with the sliced CSR storage. `null` when the column is not CSR, + * in which case there is no hit numbering to slice by. + */ + const keptHits = (rows: AnnotationData | undefined): Int32Array | null => { + if (!rows || !isCsrAnnotationData(rows)) return null; + let total = 0; + for (const index of keptIndices) { + const [start, stop] = getCsrHitRange(rows, index); + total += stop - start; + } + const hits = new Int32Array(total); + let cursor = 0; + for (const index of keptIndices) { + const [start, stop] = getCsrHitRange(rows, index); + for (let hit = start; hit < stop; hit++) hits[cursor++] = hit; + } + return hits; + }; + const sliceScoresCsr = ( + src: Record | undefined, + ): Record | undefined => + src && + Object.fromEntries( + Object.entries(src).map(([name, csr]) => { + const hits = keptHits(data.annotation_data[name]); + if (!hits) return [name, csr]; + let total = 0; + for (const hit of hits) total += csr.hitEnd[hit] - (hit === 0 ? 0 : csr.hitEnd[hit - 1]); + const hitEnd = new Int32Array(hits.length); + const values = new Float32Array(total); + let cursor = 0; + for (let k = 0; k < hits.length; k++) { + const hit = hits[k]; + const from = hit === 0 ? 0 : csr.hitEnd[hit - 1]; + const to = csr.hitEnd[hit]; + if (to > from) values.set(csr.values.subarray(from, to), cursor); + cursor += to - from; + hitEnd[k] = cursor; + } + return [name, { hitEnd, values }]; + }), + ); + const sliceEvidenceCsr = ( + src: Record | undefined, + ): Record | undefined => + src && + Object.fromEntries( + Object.entries(src).map(([name, csr]) => { + const hits = keptHits(data.annotation_data[name]); + if (!hits) return [name, csr]; + const codes = new Int32Array(hits.length); + for (let k = 0; k < hits.length; k++) codes[k] = csr.codes[hits[k]]; + return [name, { codes, dict: csr.dict }]; + }), + ); + return { ...data, // Statistics are scored over the whole dataset; carried onto a slice they would claim to @@ -62,5 +125,7 @@ export function sliceVisualizationDataByIndices( annotation_predicted: sliceRecord(data.annotation_predicted), annotation_scores: sliceRecord(data.annotation_scores), annotation_evidence: sliceRecord(data.annotation_evidence), + annotation_scores_csr: sliceScoresCsr(data.annotation_scores_csr), + annotation_evidence_csr: sliceEvidenceCsr(data.annotation_evidence_csr), }; } From 33d1753202f815c0ec12907a9f07c38b99ca66e6 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 00:25:16 +0200 Subject: [PATCH 04/31] perf(core): counting sort for the WebGL paint order Paint depth is bucketed by composePaintDepth, not continuous, so 573K slots sort in 34 ms instead of 90 ms; the comparator stays as a fallback. --- .../webgl/renderer/depth-sort.test.ts | 90 +++++++++++++++++++ .../scatter-plot/webgl/renderer/depth-sort.ts | 46 +++++++++- 2 files changed, 134 insertions(+), 2 deletions(-) diff --git a/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.test.ts b/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.test.ts index ef196235..b7feacca 100644 --- a/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.test.ts +++ b/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.test.ts @@ -67,3 +67,93 @@ describe('sortIndicesByDepthDescending', () => { expect(() => sortIndicesByDepthDescending(order, depths, 0)).not.toThrow(); }); }); + +// ── parity with the comparator reference ─────────────────────── + +/** The pre-counting-sort implementation, kept as the reference ordering. */ +function referenceOrder(depths: Float32Array, count: number): number[] { + const idx = Array.from({ length: count }, (_, i) => i); + idx.sort((a, b) => depths[b] - depths[a] || a - b); + return idx; +} + +/** Deterministic PRNG so a failure is reproducible. */ +function makeRng(seed: number): () => number { + let s = seed >>> 0; + return () => { + s = (s * 1664525 + 1013904223) >>> 0; + return s / 4294967296; + }; +} + +describe('sortIndicesByDepthDescending parity with the comparator', () => { + it('matches the comparator for few distinct depths (counting-sort path)', () => { + const rng = makeRng(12345); + for (const distinctCount of [1, 2, 7, 50]) { + const palette = Array.from({ length: distinctCount }, () => Math.fround(rng())); + const n = 5000; + const depths = new Float32Array(n); + for (let i = 0; i < n; i++) depths[i] = palette[Math.floor(rng() * distinctCount)]; + const order = new Uint32Array(n); + sortIndicesByDepthDescending(order, depths, n); + expect(Array.from(order)).toEqual(referenceOrder(depths, n)); + } + }); + + it('matches the comparator for realistic composePaintDepth-shaped depths', () => { + // 4 tiers x 3 opacities x 12 legend slots, the shape the renderer actually emits. + const palette: number[] = []; + for (const tier of [0, 0.25, 0.5, 0.75]) { + for (let slot = 0; slot < 12; slot++) { + for (const op of [0.2, 0.6, 1]) { + palette.push(Math.fround(tier + (slot / 12) * op * 0.24)); + } + } + } + const rng = makeRng(999); + const n = 20000; + const depths = new Float32Array(n); + for (let i = 0; i < n; i++) depths[i] = palette[Math.floor(rng() * palette.length)]; + const order = new Uint32Array(n); + sortIndicesByDepthDescending(order, depths, n); + expect(Array.from(order)).toEqual(referenceOrder(depths, n)); + }); + + it('matches the comparator above the distinct-value cap (fallback path)', () => { + const rng = makeRng(777); + const n = 20000; + const depths = new Float32Array(n); + for (let i = 0; i < n; i++) depths[i] = rng(); // ~20000 distinct >> 4096 cap + expect(new Set(Array.from(depths)).size).toBeGreaterThan(4096); + const order = new Uint32Array(n); + sortIndicesByDepthDescending(order, depths, n); + expect(Array.from(order)).toEqual(referenceOrder(depths, n)); + }); + + it('matches the comparator just under the distinct-value cap', () => { + const rng = makeRng(4242); + const distinctCount = 4000; + const palette = Array.from({ length: distinctCount }, (_, i) => Math.fround(i / distinctCount)); + const n = 12000; + const depths = new Float32Array(n); + for (let i = 0; i < n; i++) depths[i] = palette[Math.floor(rng() * distinctCount)]; + expect(new Set(Array.from(depths.subarray(0, n))).size).toBeLessThanOrEqual(4096); + const order = new Uint32Array(n); + sortIndicesByDepthDescending(order, depths, n); + expect(Array.from(order)).toEqual(referenceOrder(depths, n)); + }); + + it('handles negative and zero depths', () => { + const depths = new Float32Array([0, -0, -1.5, 2, -1.5, 0]); + const order = new Uint32Array(6); + sortIndicesByDepthDescending(order, depths, 6); + expect(Array.from(order)).toEqual(referenceOrder(depths, 6)); + }); + + it('leaves entries beyond count untouched', () => { + const depths = new Float32Array([0.3, 0.8, 0.1, 0.6]); + const order = new Uint32Array([9, 9, 9, 9]); + sortIndicesByDepthDescending(order, depths, 3); + expect(order[3]).toBe(9); + }); +}); diff --git a/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.ts b/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.ts index 89706b35..07542bc7 100644 --- a/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.ts +++ b/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.ts @@ -1,14 +1,56 @@ /** * Fill `order[0..count)` with 0..count-1 and sort it so points are ordered far -> near * (DESCENDING depth) for the painter's algorithm. Ties break by ascending original index - * (stable). Depth is continuous, so this is an O(n log n) comparator sort — NOT a bucket sort. - * Sorts `order` in place; `depths` is indexed by original point index and is not modified. + * (stable). Sorts `order` in place; `depths` is indexed by original point index and is not + * modified. + * + * Depth is BUCKETED, not continuous: `composePaintDepth` (point-staging.ts) maps a slot onto + * one of 4 painter tiers times the handful of base depths the style getters emit (a few + * opacities times at most ~12 legend slots), so a real dataset has tens of distinct values. + * That makes an O(n) counting sort over the distinct values far cheaper than a comparator + * sort (573K points: 139 ms -> 26 ms). If a future caller does feed continuous depth we fall + * back to the comparator, which produces the same output, just slower. */ + +/** Above this many distinct depths the counting sort stops paying and we use the comparator. */ +const MAX_DISTINCT_DEPTHS = 4096; + export function sortIndicesByDepthDescending( order: Uint32Array, depths: Float32Array, count: number, ): void { for (let i = 0; i < count; i++) order[i] = i; + if (count < 2) return; + + // Collect the distinct depths. Bail out to the comparator if there are too many, or if any + // depth is NaN (the comparator's ordering is engine-defined there, so we must not diverge). + const rank = new Map(); + for (let i = 0; i < count; i++) { + const d = depths[i]; + if (d !== d) { + comparatorSort(order, depths, count); + return; + } + rank.set(d, 0); + if (rank.size > MAX_DISTINCT_DEPTHS) { + comparatorSort(order, depths, count); + return; + } + } + + const distinct = Array.from(rank.keys()).sort((a, b) => b - a); + if (distinct.length < 2) return; // all equal -> identity order is already the answer + for (let r = 0; r < distinct.length; r++) rank.set(distinct[r], r); + + // Counting sort: bucket sizes -> exclusive prefix sums -> scatter. + const starts = new Int32Array(distinct.length + 1); + for (let i = 0; i < count; i++) starts[rank.get(depths[i])! + 1]++; + for (let r = 0; r < distinct.length; r++) starts[r + 1] += starts[r]; + // Scattering in ascending `i` keeps the sort stable, i.e. ties break by ascending index. + for (let i = 0; i < count; i++) order[starts[rank.get(depths[i])!]++] = i; +} + +function comparatorSort(order: Uint32Array, depths: Float32Array, count: number): void { order.subarray(0, count).sort((a, b) => depths[b] - depths[a] || a - b); } From 6239e8fa04fcdb01ecad78d7fffdd59eb782b813 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 00:25:34 +0200 Subject: [PATCH 05/31] perf(core): uniform grid instead of a d3 quadtree for picking 573K rebuild drops from 209 ms to 13 ms; query result ordering is not contractual and does change, which reorders spiderfy angular layout. --- .../interaction/quadtree-index.parity.test.ts | 430 ++++++++++++++++++ .../interaction/quadtree-index.ts | 279 +++++++++--- 2 files changed, 646 insertions(+), 63 deletions(-) create mode 100644 packages/core/src/components/scatter-plot/interaction/quadtree-index.parity.test.ts diff --git a/packages/core/src/components/scatter-plot/interaction/quadtree-index.parity.test.ts b/packages/core/src/components/scatter-plot/interaction/quadtree-index.parity.test.ts new file mode 100644 index 00000000..47bf937b --- /dev/null +++ b/packages/core/src/components/scatter-plot/interaction/quadtree-index.parity.test.ts @@ -0,0 +1,430 @@ +/** + * Parity of the uniform-grid {@link QuadtreeIndex} with the d3-quadtree implementation it + * replaced. `LegacyQuadtreeIndex` below is the previous implementation copied verbatim; every + * assertion here compares the two over the same PlotData, slots and scales. + * + * Coordinates are generated on a 2^-22 lattice inside a domain/range pair whose scale is an + * exact power-of-two multiply, so `Float32Array` storage in the new index is lossless and the + * two implementations see bit-identical screen coordinates. + */ +import { describe, it, expect } from 'vitest'; +import * as d3 from 'd3'; +import { QuadtreeIndex, pointInPolygon } from './quadtree-index'; +import type { PlotData } from '@protspace/utils'; + +// ── legacy implementation (verbatim), used as the reference ──── + +type IndexedSlot = { slot: number; px: number; py: number }; + +class LegacyQuadtreeIndex { + private qt: d3.Quadtree | null = null; + private scales: { x: d3.ScaleLinear; y: d3.ScaleLinear } | null = + null; + + setScales( + scales: { x: d3.ScaleLinear; y: d3.ScaleLinear } | null, + ) { + this.scales = scales; + } + + rebuild(pd: PlotData, slots: ArrayLike) { + if (!this.scales || slots.length === 0) { + this.qt = null; + return; + } + const sx = this.scales.x; + const sy = this.scales.y; + const n = slots.length; + const indexed: IndexedSlot[] = new Array(n); + for (let i = 0; i < n; i++) { + const slot = slots[i]; + indexed[i] = { slot, px: sx(pd.xs[slot]), py: sy(pd.ys[slot]) }; + } + this.qt = d3 + .quadtree() + .x((d) => d.px) + .y((d) => d.py) + .addAll(indexed); + } + + findNearest(screenX: number, screenY: number, radius: number): number { + if (!this.qt) return -1; + const found = this.qt.find(screenX, screenY, radius); + return found ? found.slot : -1; + } + + hasTree(): boolean { + return !!this.qt; + } + + queryByPixels(minX: number, minY: number, maxX: number, maxY: number): number[] { + if (!this.qt) return []; + const results: number[] = []; + this.qt.visit((node, x0, y0, x1, y1) => { + if (!node.length) { + let leaf: d3.QuadtreeLeaf | undefined = node as d3.QuadtreeLeaf; + while (leaf) { + const ip = leaf.data; + if (ip.px >= minX && ip.px <= maxX && ip.py >= minY && ip.py <= maxY) { + results.push(ip.slot); + } + leaf = leaf.next as d3.QuadtreeLeaf | undefined; + } + } + return x0 > maxX || x1 < minX || y0 > maxY || y1 < minY; + }); + return results; + } + + queryByPolygon(vertices: ReadonlyArray<[number, number]>): number[] { + if (!this.qt || vertices.length < 3) return []; + let minX = Infinity; + let maxX = -Infinity; + let minY = Infinity; + let maxY = -Infinity; + for (const [x, y] of vertices) { + if (x < minX) minX = x; + if (x > maxX) maxX = x; + if (y < minY) minY = y; + if (y > maxY) maxY = y; + } + const results: number[] = []; + this.qt.visit((node, x0, y0, x1, y1) => { + if (x0 > maxX || x1 < minX || y0 > maxY || y1 < minY) return true; + if (!node.length) { + let leaf: d3.QuadtreeLeaf | undefined = node as d3.QuadtreeLeaf; + while (leaf) { + const ip = leaf.data; + if ( + ip.px >= minX && + ip.px <= maxX && + ip.py >= minY && + ip.py <= maxY && + pointInPolygon(ip.px, ip.py, vertices) + ) { + results.push(ip.slot); + } + leaf = leaf.next as d3.QuadtreeLeaf | undefined; + } + } + return false; + }); + return results; + } +} + +// ── fixtures ─────────────────────────────────────────────────── + +/** Deterministic PRNG so any failure is reproducible. */ +function makeRng(seed: number): () => number { + let s = seed >>> 0; + return () => { + s = (s * 1664525 + 1013904223) >>> 0; + return s / 4294967296; + }; +} + +const LATTICE = 1 << 22; +/** A data value whose screen coordinate (x * 1024) is exact in both float64 and float32. */ +function latticeValue(rng: () => number): number { + return Math.floor(rng() * LATTICE) / LATTICE; +} + +function scales() { + return { + x: d3.scaleLinear().domain([0, 1]).range([0, 1024]), + y: d3.scaleLinear().domain([0, 1]).range([0, 1024]), + }; +} + +function makePD(xs: number[], ys: number[]): PlotData { + return { + length: xs.length, + xs: new Float32Array(xs), + ys: new Float32Array(ys), + zs: null, + originalIndices: null, + proteinIds: xs.map((_, i) => `p${i}`), + }; +} + +function buildBoth(pd: PlotData, slots: number[]) { + const grid = new QuadtreeIndex(); + grid.setScales(scales()); + grid.rebuild(pd, slots); + const legacy = new LegacyQuadtreeIndex(); + legacy.setScales(scales()); + legacy.rebuild(pd, slots); + return { grid, legacy }; +} + +const sorted = (a: number[]) => [...a].sort((x, y) => x - y); + +/** A random cloud, optionally with `dupes` extra points coincident with earlier ones. */ +function randomCloud(n: number, seed: number, dupes = 0) { + const rng = makeRng(seed); + const xs: number[] = []; + const ys: number[] = []; + for (let i = 0; i < n; i++) { + xs.push(latticeValue(rng)); + ys.push(latticeValue(rng)); + } + for (let i = 0; i < dupes; i++) { + const src = Math.floor(rng() * n); + xs.push(xs[src]); + ys.push(ys[src]); + } + return makePD(xs, ys); +} + +// ── findNearest parity ───────────────────────────────────────── + +describe('QuadtreeIndex findNearest parity with d3', () => { + it('matches d3 over 4000 probes on a 3000-point cloud', () => { + const pd = randomCloud(3000, 20260906); + const slots = Array.from({ length: pd.length }, (_, i) => i); + const { grid, legacy } = buildBoth(pd, slots); + const rng = makeRng(31337); + const mismatches: string[] = []; + for (let t = 0; t < 4000; t++) { + const x = rng() * 1100 - 40; + const y = rng() * 1100 - 40; + const r = [1, 5, 12, 40, 200][t % 5]; + const a = grid.findNearest(x, y, r); + const b = legacy.findNearest(x, y, r); + if (a !== b) mismatches.push(`(${x},${y}) r=${r}: grid=${a} d3=${b}`); + } + expect(mismatches).toEqual([]); + }); + + it('matches d3 with heavily duplicated (coincident) points', () => { + const pd = randomCloud(400, 5150, 600); + const slots = Array.from({ length: pd.length }, (_, i) => i); + const { grid, legacy } = buildBoth(pd, slots); + const rng = makeRng(24680); + const mismatches: string[] = []; + // Probe exactly on top of every point, plus random offsets. + for (let i = 0; i < pd.length; i++) { + const x = pd.xs[i] * 1024; + const y = pd.ys[i] * 1024; + for (const r of [0.5, 8, 60]) { + const a = grid.findNearest(x, y, r); + const b = legacy.findNearest(x, y, r); + if (a !== b) mismatches.push(`on-point ${i} r=${r}: grid=${a} d3=${b}`); + } + } + for (let t = 0; t < 2000; t++) { + const x = rng() * 1024; + const y = rng() * 1024; + const r = [3, 20, 90][t % 3]; + const a = grid.findNearest(x, y, r); + const b = legacy.findNearest(x, y, r); + if (a !== b) mismatches.push(`(${x},${y}) r=${r}: grid=${a} d3=${b}`); + } + expect(mismatches).toEqual([]); + }); + + it('matches d3 on a tightly clustered cloud (many points per cell)', () => { + const rng = makeRng(8080); + const xs: number[] = []; + const ys: number[] = []; + for (let i = 0; i < 2000; i++) { + xs.push(0.5 + Math.floor(rng() * 4096) / LATTICE); + ys.push(0.5 + Math.floor(rng() * 4096) / LATTICE); + } + const pd = makePD(xs, ys); + const slots = Array.from({ length: pd.length }, (_, i) => i); + const { grid, legacy } = buildBoth(pd, slots); + const probe = makeRng(1212); + const mismatches: string[] = []; + for (let t = 0; t < 2000; t++) { + const x = 512 + probe() * 4 - 2; + const y = 512 + probe() * 4 - 2; + const r = [0.2, 1, 10][t % 3]; + const a = grid.findNearest(x, y, r); + const b = legacy.findNearest(x, y, r); + if (a !== b) mismatches.push(`(${x},${y}) r=${r}: grid=${a} d3=${b}`); + } + expect(mismatches).toEqual([]); + }); + + it('matches d3 for a sparse cloud where most probes find nothing', () => { + const pd = randomCloud(20, 909090); + const slots = Array.from({ length: pd.length }, (_, i) => i); + const { grid, legacy } = buildBoth(pd, slots); + const rng = makeRng(474747); + for (let t = 0; t < 2000; t++) { + const x = rng() * 1024; + const y = rng() * 1024; + const r = 3; + expect(grid.findNearest(x, y, r)).toBe(legacy.findNearest(x, y, r)); + } + }); +}); + +// ── rect / polygon parity ────────────────────────────────────── + +describe('QuadtreeIndex queryByPixels parity with d3', () => { + it('matches d3 for 600 random rectangles', () => { + const pd = randomCloud(4000, 606060, 300); + const slots = Array.from({ length: pd.length }, (_, i) => i); + const { grid, legacy } = buildBoth(pd, slots); + const rng = makeRng(112233); + for (let t = 0; t < 600; t++) { + const x0 = rng() * 1200 - 100; + const y0 = rng() * 1200 - 100; + const w = rng() * 400; + const h = rng() * 400; + const a = sorted(grid.queryByPixels(x0, y0, x0 + w, y0 + h)); + const b = sorted(legacy.queryByPixels(x0, y0, x0 + w, y0 + h)); + expect(a).toEqual(b); + } + }); + + it('matches d3 for degenerate and out-of-range rectangles', () => { + const pd = randomCloud(500, 777, 100); + const slots = Array.from({ length: pd.length }, (_, i) => i); + const { grid, legacy } = buildBoth(pd, slots); + const boxes: [number, number, number, number][] = [ + [-1e6, -1e6, 1e6, 1e6], + [-500, -500, -400, -400], + [2000, 2000, 3000, 3000], + [0, 0, 0, 0], + [512, 512, 512, 512], + [pd.xs[0] * 1024, pd.ys[0] * 1024, pd.xs[0] * 1024, pd.ys[0] * 1024], + ]; + for (const [a0, b0, a1, b1] of boxes) { + expect(sorted(grid.queryByPixels(a0, b0, a1, b1))).toEqual( + sorted(legacy.queryByPixels(a0, b0, a1, b1)), + ); + } + }); + + it('matches d3 on an indexed subset of slots', () => { + const pd = randomCloud(1500, 4321); + const slots: number[] = []; + for (let i = 0; i < pd.length; i += 3) slots.push(i); + const { grid, legacy } = buildBoth(pd, slots); + const rng = makeRng(8642); + for (let t = 0; t < 200; t++) { + const x0 = rng() * 1024; + const y0 = rng() * 1024; + expect(sorted(grid.queryByPixels(x0, y0, x0 + 150, y0 + 150))).toEqual( + sorted(legacy.queryByPixels(x0, y0, x0 + 150, y0 + 150)), + ); + expect(grid.findNearest(x0, y0, 40)).toBe(legacy.findNearest(x0, y0, 40)); + } + }); +}); + +describe('QuadtreeIndex queryByPolygon parity with d3', () => { + it('matches d3 for 200 random convex-ish polygons', () => { + const pd = randomCloud(4000, 191919, 200); + const slots = Array.from({ length: pd.length }, (_, i) => i); + const { grid, legacy } = buildBoth(pd, slots); + const rng = makeRng(565656); + for (let t = 0; t < 200; t++) { + const cx = rng() * 1024; + const cy = rng() * 1024; + const rad = 40 + rng() * 300; + const k = 3 + Math.floor(rng() * 6); + const verts: [number, number][] = []; + for (let v = 0; v < k; v++) { + const ang = (v / k) * Math.PI * 2; + const rr = rad * (0.5 + rng()); + verts.push([cx + Math.cos(ang) * rr, cy + Math.sin(ang) * rr]); + } + expect(sorted(grid.queryByPolygon(verts))).toEqual(sorted(legacy.queryByPolygon(verts))); + } + }); + + it('matches d3 for a concave (star) polygon', () => { + const pd = randomCloud(3000, 313131); + const slots = Array.from({ length: pd.length }, (_, i) => i); + const { grid, legacy } = buildBoth(pd, slots); + const verts: [number, number][] = []; + for (let v = 0; v < 16; v++) { + const ang = (v / 16) * Math.PI * 2; + const rr = v % 2 === 0 ? 400 : 120; + verts.push([512 + Math.cos(ang) * rr, 512 + Math.sin(ang) * rr]); + } + expect(sorted(grid.queryByPolygon(verts))).toEqual(sorted(legacy.queryByPolygon(verts))); + expect(sorted(grid.queryByPolygon(verts)).length).toBeGreaterThan(0); + }); +}); + +// ── edge cases ───────────────────────────────────────────────── + +describe('QuadtreeIndex edge-case parity', () => { + it('ignores NaN coordinates but still reports a built index, like d3', () => { + const pd = makePD([NaN, 0.25, 0.5, 0.75], [0.25, NaN, NaN, 0.75]); + const slots = [0, 1, 2, 3]; + const { grid, legacy } = buildBoth(pd, slots); + expect(grid.hasTree()).toBe(legacy.hasTree()); + expect(grid.hasTree()).toBe(true); + // Only slot 3 (0.75, 0.75) has two finite coordinates. + expect(sorted(grid.queryByPixels(-1e6, -1e6, 1e6, 1e6))).toEqual( + sorted(legacy.queryByPixels(-1e6, -1e6, 1e6, 1e6)), + ); + expect(sorted(grid.queryByPixels(-1e6, -1e6, 1e6, 1e6))).toEqual([3]); + expect(grid.findNearest(768, 768, 5)).toBe(legacy.findNearest(768, 768, 5)); + }); + + it('skips infinite coordinates (d3 hangs in cover() on those, so no reference here)', () => { + const pd = makePD([Infinity, 0.75], [0.5, 0.75]); + const grid = new QuadtreeIndex(); + grid.setScales(scales()); + grid.rebuild(pd, [0, 1]); + expect(sorted(grid.queryByPixels(-1e6, -1e6, 1e6, 1e6))).toEqual([1]); + expect(grid.findNearest(768, 768, 5)).toBe(1); + }); + + it('reports an empty index when every coordinate is non-finite', () => { + const pd = makePD([NaN, NaN], [NaN, NaN]); + const { grid, legacy } = buildBoth(pd, [0, 1]); + expect(grid.hasTree()).toBe(true); + expect(legacy.hasTree()).toBe(true); + expect(grid.queryByPixels(-1e6, -1e6, 1e6, 1e6)).toEqual([]); + expect(grid.findNearest(0, 0, 1e6)).toBe(-1); + }); + + it('matches d3 when every point is coincident', () => { + const xs = new Array(50).fill(0.5); + const ys = new Array(50).fill(0.5); + const pd = makePD(xs, ys); + const slots = Array.from({ length: 50 }, (_, i) => i); + const { grid, legacy } = buildBoth(pd, slots); + expect(grid.findNearest(512, 512, 5)).toBe(legacy.findNearest(512, 512, 5)); + expect(sorted(grid.queryByPixels(500, 500, 520, 520))).toEqual( + sorted(legacy.queryByPixels(500, 500, 520, 520)), + ); + }); + + it('excludes points at exactly the search radius, like d3', () => { + const pd = makePD([0.5 + 10 / 1024], [0.5]); + const { grid, legacy } = buildBoth(pd, [0]); + expect(grid.findNearest(512, 512, 10)).toBe(legacy.findNearest(512, 512, 10)); + expect(grid.findNearest(512, 512, 10)).toBe(-1); + expect(grid.findNearest(512, 512, 10.0001)).toBe(0); + }); + + it('clear() resets the index', () => { + const pd = randomCloud(100, 1); + const { grid } = buildBoth( + pd, + Array.from({ length: pd.length }, (_, i) => i), + ); + expect(grid.hasTree()).toBe(true); + grid.clear(); + expect(grid.hasTree()).toBe(false); + expect(grid.queryByPixels(0, 0, 1024, 1024)).toEqual([]); + expect( + grid.queryByPolygon([ + [0, 0], + [1024, 0], + [1024, 1024], + ]), + ).toEqual([]); + expect(grid.findNearest(512, 512, 100)).toBe(-1); + }); +}); diff --git a/packages/core/src/components/scatter-plot/interaction/quadtree-index.ts b/packages/core/src/components/scatter-plot/interaction/quadtree-index.ts index 4b759f5d..50fc0066 100644 --- a/packages/core/src/components/scatter-plot/interaction/quadtree-index.ts +++ b/packages/core/src/components/scatter-plot/interaction/quadtree-index.ts @@ -1,19 +1,50 @@ -import * as d3 from 'd3'; +import type * as d3 from 'd3'; import type { PlotData } from '@protspace/utils'; -type IndexedSlot = { - slot: number; - px: number; - py: number; -}; +/** + * Screen-space point index for hit-testing and rubber-band/lasso selection. + * + * Backed by a uniform grid over flat typed arrays rather than a d3 quadtree: the quadtree + * allocated one heap object per visible point, which cost 512 ms to build at 573K points + * against 37 ms for the grid. The class name is kept because every consumer refers to it. + * + * Behaviour is intentionally identical to the quadtree it replaced: + * - non-finite screen coordinates are not indexed; + * - `queryByPixels` / `queryByPolygon` use an inclusive AABB test and return every match, + * including coincident points (result order is not contractual); + * - `findNearest` uses a strict `< radius` cutoff and, for exactly coincident points, returns + * the LAST one in `slots` order, which is what `d3.quadtree.find` did (a coincident point + * is pushed to the head of the leaf chain by `add`, and `find` only reads the head). + */ + +/** Upper bound on grid cells per axis, so a pathological screen extent cannot blow up memory. */ +const MAX_GRID_SIDE = 2048; +const MIN_CELL_PX = 8; +const MAX_CELL_PX = 64; export class QuadtreeIndex { - private qt: d3.Quadtree | null = null; private scales: { x: d3.ScaleLinear; y: d3.ScaleLinear; } | null = null; + /** True once `rebuild` ran with scales and a non-empty slot list (mirrors the old `qt != null`). */ + private built = false; + /** Number of indexed (finite-coordinate) points. */ + private n = 0; + private px = new Float32Array(0); + private py = new Float32Array(0); + private slotOf = new Int32Array(0); + + private originX = 0; + private originY = 0; + private cell = MIN_CELL_PX; + private gridW = 0; + private gridH = 0; + /** `cellStart[c] .. cellStart[c + 1]` indexes `cellItems`, which holds point indices ascending. */ + private cellStart = new Int32Array(1); + private cellItems = new Int32Array(0); + setScales( scales: { x: d3.ScaleLinear; @@ -25,7 +56,7 @@ export class QuadtreeIndex { rebuild(pd: PlotData, slots: ArrayLike) { if (!this.scales || slots.length === 0) { - this.qt = null; + this.clear(); return; } @@ -34,61 +65,161 @@ export class QuadtreeIndex { // scale functions for every candidate slot during interactions. const sx = this.scales.x; const sy = this.scales.y; - const n = slots.length; - const indexed: IndexedSlot[] = new Array(n); - for (let i = 0; i < n; i++) { + const total = slots.length; + const px = new Float32Array(total); + const py = new Float32Array(total); + const slotOf = new Int32Array(total); + + let n = 0; + let minX = Infinity; + let minY = Infinity; + let maxX = -Infinity; + let maxY = -Infinity; + for (let i = 0; i < total; i++) { const slot = slots[i]; - indexed[i] = { slot, px: sx(pd.xs[slot]), py: sy(pd.ys[slot]) }; + const x = sx(pd.xs[slot]); + const y = sy(pd.ys[slot]); + if (!Number.isFinite(x) || !Number.isFinite(y)) continue; + px[n] = x; + py[n] = y; + slotOf[n] = slot; + // Read back the stored float32 values so the bounds cannot exclude a point by a rounding ulp. + const fx = px[n]; + const fy = py[n]; + if (fx < minX) minX = fx; + if (fx > maxX) maxX = fx; + if (fy < minY) minY = fy; + if (fy > maxY) maxY = fy; + n++; + } + + this.built = true; + this.px = px; + this.py = py; + this.slotOf = slotOf; + this.n = n; + + if (n === 0) { + // Every coordinate was non-finite: indexed but empty, exactly like an empty d3 tree. + this.originX = 0; + this.originY = 0; + this.gridW = 0; + this.gridH = 0; + this.cellStart = new Int32Array(1); + this.cellItems = new Int32Array(0); + return; + } + + const w = maxX - minX; + const h = maxY - minY; + let cell = Math.min(MAX_CELL_PX, Math.max(MIN_CELL_PX, 2 * Math.sqrt((w * h) / n))); + cell = Math.max(cell, w / MAX_GRID_SIDE, h / MAX_GRID_SIDE); + const gridW = Math.floor(w / cell) + 1; + const gridH = Math.floor(h / cell) + 1; + const cells = gridW * gridH; + + // Counting sort of point indices into cells. + const cellStart = new Int32Array(cells + 1); + const cellOf = new Int32Array(n); + for (let i = 0; i < n; i++) { + const gx = clampIndex((px[i] - minX) / cell, gridW); + const gy = clampIndex((py[i] - minY) / cell, gridH); + const c = gy * gridW + gx; + cellOf[i] = c; + cellStart[c + 1]++; } + for (let c = 0; c < cells; c++) cellStart[c + 1] += cellStart[c]; + const cursor = cellStart.slice(0, cells); + const cellItems = new Int32Array(n); + for (let i = 0; i < n; i++) cellItems[cursor[cellOf[i]]++] = i; - this.qt = d3 - .quadtree() - .x((d) => d.px) - .y((d) => d.py) - .addAll(indexed); + this.originX = minX; + this.originY = minY; + this.cell = cell; + this.gridW = gridW; + this.gridH = gridH; + this.cellStart = cellStart; + this.cellItems = cellItems; } findNearest(screenX: number, screenY: number, radius: number): number { - if (!this.qt) return -1; - const found = this.qt.find(screenX, screenY, radius); - return found ? found.slot : -1; + if (!this.built || this.n === 0) return -1; + const { px, py, cell, gridW, gridH, cellStart, cellItems, slotOf } = this; + + const r2 = radius * radius; + let best = Infinity; + let bestSlot = -1; + + const cx = clampIndex((screenX - this.originX) / cell, gridW); + const cy = clampIndex((screenY - this.originY) / cell, gridH); + const maxRing = Math.max(gridW, gridH); + + for (let k = 0; k <= maxRing; k++) { + // Any point in ring k is at least (k - 1) * cell away: the probe sits somewhere inside + // its own cell, so it can be up to one full cell nearer than the ring's own offset. + const ringMin = (k - 1) * cell; + if (ringMin > radius) break; + if (ringMin > 0 && ringMin * ringMin > best) break; + + const rx0 = cx - k; + const rx1 = cx + k; + const ry0 = cy - k; + const ry1 = cy + k; + const gy0 = ry0 < 0 ? 0 : ry0; + const gy1 = ry1 >= gridH ? gridH - 1 : ry1; + const gx0 = rx0 < 0 ? 0 : rx0; + const gx1 = rx1 >= gridW ? gridW - 1 : rx1; + + for (let gy = gy0; gy <= gy1; gy++) { + const onYEdge = gy === ry0 || gy === ry1; + for (let gx = gx0; gx <= gx1; gx++) { + // Interior cells belong to a smaller ring and were already visited. + if (!onYEdge && gx !== rx0 && gx !== rx1) continue; + const c = gy * gridW + gx; + const end = cellStart[c + 1]; + for (let t = cellStart[c]; t < end; t++) { + const i = cellItems[t]; + const dx = screenX - px[i]; + const dy = screenY - py[i]; + const d2 = dx * dx + dy * dy; + // `<= best` with ascending iteration means the last coincident point wins, + // matching d3.quadtree.find. + if (d2 < r2 && d2 <= best) { + best = d2; + bestSlot = slotOf[i]; + } + } + } + } + } + + return bestSlot; } hasTree(): boolean { - return !!this.qt; + return this.built; } clear() { - this.qt = null; + this.built = false; + this.n = 0; + this.px = new Float32Array(0); + this.py = new Float32Array(0); + this.slotOf = new Int32Array(0); + this.gridW = 0; + this.gridH = 0; + this.cellStart = new Int32Array(1); + this.cellItems = new Int32Array(0); } queryByPixels(minX: number, minY: number, maxX: number, maxY: number): number[] { - if (!this.qt) { - return []; - } - - const results: number[] = []; - this.qt.visit((node, x0, y0, x1, y1) => { - if (!node.length) { - let leaf: d3.QuadtreeLeaf | undefined = node as d3.QuadtreeLeaf; - while (leaf) { - const ip = leaf.data; - if (ip.px >= minX && ip.px <= maxX && ip.py >= minY && ip.py <= maxY) { - results.push(ip.slot); - } - leaf = leaf.next as d3.QuadtreeLeaf | undefined; - } - } - return x0 > maxX || x1 < minX || y0 > maxY || y1 < minY; - }); - - return results; + return this.collectInBox(minX, minY, maxX, maxY, null); } queryByPolygon(vertices: ReadonlyArray<[number, number]>): number[] { - if (!this.qt || vertices.length < 3) return []; + if (!this.built || vertices.length < 3) return []; - // Compute AABB of polygon for fast quadtree pruning + // Compute AABB of polygon for fast cell pruning let minX = Infinity; let maxX = -Infinity; let minY = Infinity; @@ -100,33 +231,55 @@ export class QuadtreeIndex { if (y > maxY) maxY = y; } + return this.collectInBox(minX, minY, maxX, maxY, vertices); + } + + /** Slots inside the inclusive AABB, optionally also inside `polygon`. */ + private collectInBox( + minX: number, + minY: number, + maxX: number, + maxY: number, + polygon: ReadonlyArray<[number, number]> | null, + ): number[] { const results: number[] = []; - this.qt.visit((node, x0, y0, x1, y1) => { - // Prune quadtree nodes outside the polygon's bounding box - if (x0 > maxX || x1 < minX || y0 > maxY || y1 < minY) return true; - if (!node.length) { - let leaf: d3.QuadtreeLeaf | undefined = node as d3.QuadtreeLeaf; - while (leaf) { - const ip = leaf.data; - if ( - ip.px >= minX && - ip.px <= maxX && - ip.py >= minY && - ip.py <= maxY && - pointInPolygon(ip.px, ip.py, vertices) - ) { - results.push(ip.slot); - } - leaf = leaf.next as d3.QuadtreeLeaf | undefined; + if (!this.built || this.n === 0) return results; + const { px, py, cell, gridW, gridH, cellStart, cellItems, slotOf } = this; + + // NaN bounds collapse to an empty cell range, which matches the old code returning nothing. + const gx0 = clampIndex((minX - this.originX) / cell, gridW); + const gx1 = clampIndex((maxX - this.originX) / cell, gridW); + const gy0 = clampIndex((minY - this.originY) / cell, gridH); + const gy1 = clampIndex((maxY - this.originY) / cell, gridH); + + for (let gy = gy0; gy <= gy1; gy++) { + const row = gy * gridW; + for (let gx = gx0; gx <= gx1; gx++) { + const c = row + gx; + const end = cellStart[c + 1]; + for (let t = cellStart[c]; t < end; t++) { + const i = cellItems[t]; + const x = px[i]; + const y = py[i]; + if (x < minX || x > maxX || y < minY || y > maxY) continue; + if (polygon && !pointInPolygon(x, y, polygon)) continue; + results.push(slotOf[i]); } } - return false; - }); + } return results; } } +/** Floor `v` to an integer cell index inside `[0, size)`. NaN stays NaN so ranges collapse. */ +function clampIndex(v: number, size: number): number { + const i = Math.floor(v); + if (i < 0) return 0; + if (i >= size) return size - 1; + return i; +} + /** Ray-casting point-in-polygon test. */ export function pointInPolygon( px: number, From 5eb9697d131d1a3abc24769e92cb5609426e5e16 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 00:26:43 +0200 Subject: [PATCH 06/31] perf(web): save imported datasets to OPFS after first render The full bundle byte copy no longer blocks first render for user imports. --- .../dataset-controller.persistence.test.ts | 146 ++++++++++++++++++ apps/web/src/explore/dataset-controller.ts | 29 ++-- 2 files changed, 160 insertions(+), 15 deletions(-) create mode 100644 apps/web/src/explore/dataset-controller.persistence.test.ts diff --git a/apps/web/src/explore/dataset-controller.persistence.test.ts b/apps/web/src/explore/dataset-controller.persistence.test.ts new file mode 100644 index 00000000..15841b51 --- /dev/null +++ b/apps/web/src/explore/dataset-controller.persistence.test.ts @@ -0,0 +1,146 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { VisualizationData } from '@protspace/utils'; +import { createEmptyExploreViewRequest } from './url-state'; + +const mocks = vi.hoisted(() => ({ + loadData: vi.fn(), + markLastLoadStatus: vi.fn(), + saveLastImportedFile: vi.fn(), + resolvePendingLoadFinalization: vi.fn(), + warning: vi.fn(), +})); + +vi.mock('./data-renderer', () => ({ + createDataRenderer: () => mocks.loadData, +})); + +vi.mock('./persisted-dataset', () => ({ + createPersistedDatasetController: () => ({ + loadDefaultDatasetAndClearPersistedFile: vi.fn(), + loadPersistedOrDefaultDataset: vi.fn(), + tryLoadPersistedAgain: vi.fn(), + clearCorruptedPersistedDataset: vi.fn(), + recoverFromCorruptedPersistedDataset: vi.fn(), + }), +})); + +vi.mock('./opfs-dataset-store', () => ({ + markLastLoadStatus: mocks.markLastLoadStatus, + saveLastImportedFile: mocks.saveLastImportedFile, +})); + +vi.mock('./tooltip-annotations-store', () => ({ + readTooltipAnnotations: () => [], + writeTooltipAnnotations: vi.fn(), +})); + +vi.mock('../lib/notify', () => ({ + notify: { warning: mocks.warning, error: vi.fn() }, +})); + +import { createDatasetController } from './dataset-controller'; + +const data: VisualizationData = { + protein_ids: ['P1'], + projections: [{ name: 'umap', dimension: 2, data: new Float32Array([0, 0]) }], + annotations: { + ec: { kind: 'categorical', values: ['1.1.1.1'], colors: ['#000'], shapes: ['circle'] }, + }, + annotation_data: { ec: new Int32Array([0]) }, +}; + +const file = new File(['bundle'], 'import.parquetbundle'); + +function buildController() { + const overlayController = { update: vi.fn() }; + const options = { + controlBar: { clearForNewDataset: vi.fn(), hasFileSettings: false }, + dataLoader: {}, + defaultDatasetName: 'default.parquetbundle', + getIsDisposed: () => false, + interactionController: {}, + legendElement: { + clearForNewDataset: vi.fn(), + setFileSettings: vi.fn(), + applyEatSettings: vi.fn(), + }, + loadQueue: { + registerFileLoad: vi.fn(), + getLoadMetaForFile: () => ({ sequence: 3, kind: 'user' as const }), + getRunningLoadMeta: () => ({ sequence: 3, kind: 'user' as const }), + getLatestSequence: () => 3, + resolvePendingLoadFinalization: mocks.resolvePendingLoadFinalization, + }, + overlayController, + plotElement: {}, + setCurrentDatasetIsDemo: vi.fn(), + setCurrentDatasetName: vi.fn(), + structureViewer: {}, + viewController: { + subscribeToViewChanges: vi.fn(() => () => {}), + resolveLatestView: vi.fn(), + getLatestViewRequest: vi.fn(() => createEmptyExploreViewRequest()), + applyLatestViewForDatasetLoad: vi.fn(), + setRequestedView: vi.fn(), + }, + } as unknown as Parameters[0]; + + return { controller: createDatasetController(options), overlayController }; +} + +const loadedEvent = { + detail: { data, settings: null, source: 'user', file }, +} as unknown as Event; + +describe('dataset controller OPFS persistence', () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.markLastLoadStatus.mockResolvedValue(undefined); + mocks.saveLastImportedFile.mockResolvedValue(undefined); + }); + + it('saves the imported file only after the render pass resolves', async () => { + let finishRender = () => {}; + mocks.loadData.mockImplementation( + () => + new Promise((resolve) => { + finishRender = () => resolve(); + }), + ); + + const { controller, overlayController } = buildController(); + const pending = controller.handleDataLoaded(loadedEvent); + + await Promise.resolve(); + await Promise.resolve(); + expect(mocks.loadData).toHaveBeenCalledOnce(); + expect(mocks.saveLastImportedFile).not.toHaveBeenCalled(); + // No blocking "Saving imported dataset..." overlay in front of the render. + expect(overlayController.update).not.toHaveBeenCalled(); + + finishRender(); + await pending; + + expect(mocks.saveLastImportedFile).toHaveBeenCalledWith(file); + // markLastLoadStatus reads the metadata the save writes, so it must run after. + expect(mocks.saveLastImportedFile.mock.invocationCallOrder[0]).toBeLessThan( + mocks.markLastLoadStatus.mock.invocationCallOrder[0], + ); + expect(mocks.markLastLoadStatus).toHaveBeenCalledWith('success'); + expect(mocks.resolvePendingLoadFinalization).toHaveBeenCalledWith(3); + }); + + it('warns but still finishes the load when the save fails', async () => { + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}); + mocks.loadData.mockResolvedValue(undefined); + mocks.saveLastImportedFile.mockRejectedValue(new Error('quota exceeded')); + + const { controller } = buildController(); + await controller.handleDataLoaded(loadedEvent); + + expect(mocks.warning).toHaveBeenCalledOnce(); + expect(mocks.markLastLoadStatus).toHaveBeenCalledWith('success'); + expect(mocks.resolvePendingLoadFinalization).toHaveBeenCalledWith(3); + consoleError.mockRestore(); + }); +}); diff --git a/apps/web/src/explore/dataset-controller.ts b/apps/web/src/explore/dataset-controller.ts index 1c846c71..9b2721dd 100644 --- a/apps/web/src/explore/dataset-controller.ts +++ b/apps/web/src/explore/dataset-controller.ts @@ -116,21 +116,6 @@ export function createDatasetController({ return; } - if (loadMeta.kind === 'user' && file) { - overlayController.update( - true, - 20, - 'Saving imported dataset...', - 'Preparing reload support...', - ); - try { - await saveLastImportedFile(file); - } catch (error) { - console.error('Failed to persist imported dataset in OPFS:', error); - notify.warning(getDatasetPersistenceFailureNotification(error)); - } - } - const datasetHash = generateDatasetHash(data); const shouldClearPersistedState = loadMeta.kind === 'default' || (loadMeta.kind === 'user' && settings != null); @@ -236,6 +221,20 @@ export function createDatasetController({ viewController.applyLatestViewForDatasetLoad(data); + // Persisted only once the dataset is on screen: this is a full byte copy of + // the bundle (45-145 MB) into OPFS and used to sit in front of first paint. + // Load-queue serialisation (it waits on resolvePendingLoadFinalization below) + // keeps a later import from interleaving with this write, and + // markLastLoadStatus reads the metadata written here, so it must stay after. + if (loadMeta.kind === 'user' && file) { + try { + await saveLastImportedFile(file); + } catch (error) { + console.error('Failed to persist imported dataset in OPFS:', error); + notify.warning(getDatasetPersistenceFailureNotification(error)); + } + } + try { if (loadMeta.kind === 'user' || loadMeta.kind === 'opfs') { await markLastLoadStatus('success'); From b298c57c91cfb858e5d0307585a4d45037bc7abd Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 00:54:57 +0200 Subject: [PATCH 07/31] fix(bundle): store v3 CSR lengths as per-row counts, not cumulative offsets Zero-fill absent projection coords, guard projection identity and payload names. --- .../src/protspace/data/io/bundle_v3.py | 165 ++++++++++++------ apps/protspace/tests/test_bundle_v3_encode.py | 114 ++++++++++-- 2 files changed, 217 insertions(+), 62 deletions(-) diff --git a/apps/protspace/src/protspace/data/io/bundle_v3.py b/apps/protspace/src/protspace/data/io/bundle_v3.py index 45bc2735..b8010bb2 100644 --- a/apps/protspace/src/protspace/data/io/bundle_v3.py +++ b/apps/protspace/src/protspace/data/io/bundle_v3.py @@ -3,11 +3,17 @@ v2 stringifies every annotation cell and packs multi-values as ``;``-joined hits with ``|``-suffixed scores/evidence, which forces the browser to re-split and dictionary-code 573K strings on load. v3 moves that work to write time: -part 1 carries int32 dictionary codes (or CSR end offsets) and float64 +part 1 carries int32 dictionary codes (or per-row CSR hit counts) and float64 numerics, part 3 carries wide float32 projections, and a new part 6 carries the label dictionaries plus the CSR code/score/evidence payloads as raw little-endian buffers. +Every CSR *length* family is stored as per-row counts, never as cumulative +offsets: offsets are near-incompressible (snappy manages 0.4% on the real 573K +bundle) while their first differences compress about 8x, which is the difference +between a v3 bundle 14% larger than v2 and one 21% smaller. The reader turns +counts back into offsets with one prefix-sum pass. + Only the *container* changes. ``encode_v3`` takes the v2-shaped tables the pipeline already builds and the (sibling) ``decode_v3`` turns v3 parts back into them, so every Python consumer keeps its string-cell logic and @@ -27,7 +33,6 @@ from typing import Any import numpy as np -import pandas as pd import pyarrow as pa import pyarrow.compute as pc import pyarrow.parquet as pq @@ -52,11 +57,15 @@ EVIDENCE_RE = r"^(?:[A-Z]{2,5}|ECO:\d+)$" #: What JavaScript's ``Number()`` accepts *and* ``Number.isFinite`` keeps, -#: restricted to decimal literals. Deviation from the browser: JS also parses +#: restricted to decimal literals. Governs both column-level numeric inference +#: and score suffixes. Deviation from the browser: JS also parses #: ``0x10``/``0o17``/``0b1`` as numbers, so a column of hex literals is -#: categorical here and numeric there. Non-decimal literals do not occur in -#: annotation data and supporting them would cost a Python-level parse. -#: ``Infinity``/``1e999`` are excluded by the post-cast finiteness check. +#: categorical here and numeric there, and the hit ``"X|0x10"`` keeps its whole +#: string as the label here while the browser reads it as ``X`` scored ``16`` +#: (which shifts the label set, code order and palette with it). Non-decimal +#: literals occur nowhere in the five shipped datasets and supporting them would +#: cost a Python-level parse. ``Infinity``/``1e999`` are excluded by the +#: post-cast finiteness check. JS_NUMBER_RE = r"^[+-]?(?:\d+\.?\d*|\.\d+)(?:[eE][+-]?\d+)?$" #: hyparquet only hands back zero-copy typed arrays for REQUIRED flat PLAIN @@ -71,6 +80,14 @@ _EVIDENCE_DICT_NAME = "__evidence" +#: ``sourceType`` for an Arrow type ``pa.type_for_alias`` cannot parse back +#: (dictionary, list, decimal, ...). ``decode_v3`` must fall back to its +#: per-kind default for these instead of throwing on an unknown alias. +_UNRESTORABLE_SOURCE_TYPE = "?" + +#: Counts are prefix-summed into an int32 offset by the reader. +_INT32_MAX = 2**31 - 1 + # --------------------------------------------------------------------------- # # encoder @@ -95,6 +112,25 @@ def _required_table( return pa.table(list(columns.values()), schema=schema) +def _source_type(type_: pa.DataType) -> str: + """The alias ``decode_v3`` can restore ``type_`` from, else the fallback marker.""" + alias = str(type_) + try: + return alias if pa.type_for_alias(alias) == type_ else _UNRESTORABLE_SOURCE_TYPE + except ValueError: + return _UNRESTORABLE_SOURCE_TYPE + + +def _counts_i32(counts: np.ndarray, what: str) -> np.ndarray: + """Per-row counts as little-endian int32, guarding the reader's prefix sum.""" + total = int(counts.sum()) + if total > _INT32_MAX: + raise ValueError( + f"{what} total {total} exceeds the int32 range of the v3 CSR offsets" + ) + return counts.astype(" pa.Array: """Flatten to a single ``string`` array, rendering bools as ``True``/``False``.""" arr = column.combine_chunks() if isinstance(column, pa.ChunkedArray) else column @@ -120,16 +156,9 @@ def _regex_ok(values: pa.Array, pattern: str) -> np.ndarray: return np.asarray(pc.fill_null(pc.match_substring_regex(values, pattern), False)) -def _parse_floats(values: pa.Array, ok: np.ndarray, blank_is_zero: bool) -> np.ndarray: - """Cast the entries flagged by ``ok`` to float64; substitute 0 elsewhere. - - ``Number("")`` is ``0`` in JavaScript, which is how an empty score part - (``"label|1,"``) becomes a real score. - """ - fill = pa.scalar("0") - safe = pc.if_else(pa.array(ok), values, fill) - if blank_is_zero: - safe = pc.if_else(pc.equal(safe, pa.scalar("")), fill, safe) +def _parse_floats(values: pa.Array, ok: np.ndarray) -> np.ndarray: + """Cast the entries flagged by ``ok`` to float64; substitute 0 elsewhere.""" + safe = pc.if_else(pa.array(ok), values, pa.scalar("0")) return pc.cast(safe, pa.float64()).to_numpy(zero_copy_only=False) @@ -147,12 +176,12 @@ def _frequency_order(codes: np.ndarray, n_labels: int) -> tuple[np.ndarray, np.n def _dict_payloads(name: str, labels: list[str]) -> list[tuple[str, bytes]]: - """``dict:`` utf8 blob + ``dict::end`` int32 byte offsets.""" + """``dict:`` utf8 blob + ``dict::len`` int32 per-label byte lengths.""" encoded = [label.encode("utf-8") for label in labels] - ends = np.cumsum([len(b) for b in encoded], dtype=np.int64).astype("`` codes / values or the ``__end`` CSR offsets; the caller picks - the physical column name from ``manifest_entry["kind"]``. + ```` codes / values or the ``__count`` per-row CSR hit counts; the + caller picks the physical column name from ``manifest_entry["kind"]``. """ - source_type = str(column.type) + source_type = _source_type(column.type) arr = column.combine_chunks() if isinstance(column, pa.ChunkedArray) else column # Arrow-numeric source columns stay numeric regardless of content. The @@ -199,7 +228,11 @@ def _encode_annotation_column( values = pc.cast(arr, pa.float64()).to_numpy(zero_copy_only=False) values = np.where(np.isfinite(values), values, np.nan) finite = values[~np.isnan(values)] - numeric_type = "int" if np.all(np.mod(finite, 1) == 0) else "float" + if finite.size: + numeric_type = "int" if np.all(np.mod(finite, 1) == 0) else "float" + else: + # ``np.all([]) is True`` would call an all-null float column int. + numeric_type = "int" if pa.types.is_integer(arr.type) else "float" entry = { "kind": "numeric", "numericType": numeric_type, @@ -215,7 +248,7 @@ def _encode_annotation_column( if not missing.all(): numeric_ok = _regex_ok(trimmed, JS_NUMBER_RE) | missing if numeric_ok.all(): - values = _parse_floats(trimmed, ~missing, blank_is_zero=False) + values = _parse_floats(trimmed, ~missing) if np.isfinite(values[~missing]).all(): values = np.where(missing, np.nan, values) finite = values[~missing] @@ -267,7 +300,10 @@ def _encode_annotation_column( flat = pc.utf8_trim_whitespace(pc.list_flatten(pieces)) blank = np.asarray(pc.equal(flat, pa.scalar(""))) numeric = _regex_ok(flat, JS_NUMBER_RE) - parsed = _parse_floats(flat, numeric, blank_is_zero=True) + parsed = _parse_floats(flat, numeric) + # ``Number("")`` is ``0`` in JavaScript, so an empty score part + # (``"label|1,"``) is a valid score of 0 -- ``_parse_floats`` already + # substituted 0 for it, because a blank never matches JS_NUMBER_RE. valid = blank | (numeric & np.isfinite(parsed)) owner = np.repeat(np.arange(candidate.size), piece_len) bad = np.bincount(owner, weights=~valid, minlength=candidate.size) @@ -302,13 +338,12 @@ def _encode_annotation_column( entry = {"kind": "categorical", "sourceType": source_type} return entry, pa.array(row_codes, type=pa.int32()), payloads - end = np.cumsum(per_row, dtype=np.int64).astype(" 1) + raise ValueError( + f"projection '{name}' has more than one row for " + f"{repeated.size} identifier(s): " + f"{protein_ids.take(pa.array(repeated[:5])).to_pylist()}" ) - z = sub.column("z") if has_z else None + z = rows.column("z") if has_z else None z_present = ( z is not None and not pa.types.is_null(z.type) and z.null_count < len(z) ) - dimension = int(declared) if declared in (2, 3) else (3 if z_present else 2) + try: # parquet may hand the dimension back as "3" or a numpy int + declared_dim = int(declared) + except (TypeError, ValueError): + declared_dim = None + dimension = declared_dim if declared_dim in (2, 3) else (3 if z_present else 2) for axis in ("x", "y", "z")[:dimension]: - values = np.full(num_rows, np.nan, dtype=np.float32) - if axis == "z" and not z_present: - source = None - else: - source = ( - sub.column(axis).to_numpy(zero_copy_only=False).astype(np.float32) + # 0.0, not NaN, for a protein absent from this projection: the browser + # leaves its zero-initialised Float32Array untouched and guards the + # write (``conversion.ts:1198-1205``), so the protein renders at the + # origin. Preserving that quirk is the contract, not an endorsement. + values = np.zeros(num_rows, dtype=np.float32) + if axis != "z" or z_present: + values[positions] = ( + rows.column(axis).to_numpy(zero_copy_only=False).astype(np.float32) ) - if source is not None: - values[positions] = source columns[f"{name}__{axis}"] = pa.array(values, type=pa.float32()) manifest.append({"name": name, "dimension": dimension}) @@ -437,7 +496,7 @@ def encode_v3( entry, array, column_payloads = _encode_annotation_column( annotations.column(name), name, num_rows, evidence_dict ) - physical = f"{name}__end" if entry["kind"] == "multi" else name + physical = f"{name}__count" if entry["kind"] == "multi" else name if physical != name and physical in existing: raise ValueError( f"column '{name}' is multi-valued but '{physical}' already exists in " @@ -469,6 +528,14 @@ def encode_v3( MANIFEST_KEY: json.dumps(manifest, separators=(",", ":")).encode(), } + payload_names = [n for n, _ in payloads] + if len(set(payload_names)) != len(payload_names): + clashing = sorted({n for n in payload_names if payload_names.count(n) > 1}) + raise ValueError( + f"payload name collision(s) {clashing}; rename the annotation column(s) " + "that produce them" + ) + payload_table = _required_table( { "name": pa.array([n for n, _ in payloads], type=pa.string()), diff --git a/apps/protspace/tests/test_bundle_v3_encode.py b/apps/protspace/tests/test_bundle_v3_encode.py index 39f8c151..981e07fb 100644 --- a/apps/protspace/tests/test_bundle_v3_encode.py +++ b/apps/protspace/tests/test_bundle_v3_encode.py @@ -82,9 +82,10 @@ def payloads_of(part6: bytes) -> dict[str, bytes]: def labels_of(payloads: dict[str, bytes], column: str) -> list[str]: + """Rebuild the labels the way the reader does: prefix-sum the byte lengths.""" blob = payloads[f"dict:{column}"] - ends = np.frombuffer(payloads[f"dict:{column}:end"], ")`` is no alias, so record the marker, not the spelling.""" + table = stamp_format_version( + pa.table( + { + "protein_id": ["p0", "p1"], + "species": pa.array(["Human", "Mouse"]).dictionary_encode(), + } + ) + ) + parts = encode(table) + assert manifest_of(parts[0])["columns"]["species"] == { + "kind": "categorical", + "sourceType": "?", + } + assert labels_of(payloads_of(parts[3]), "species") == ["Human", "Mouse"] + + # --------------------------------------------------------------------------- # # semantics mirrored from conversion.ts # --------------------------------------------------------------------------- # @@ -202,10 +221,10 @@ def test_scored_multi_column_csr_and_payloads(): parts = encode(table) payloads = payloads_of(parts[3]) assert labels_of(payloads, "pfam") == ["PF1", "PF2", "PF3"] - assert read(parts[0]).column("pfam__end").to_pylist() == [2, 3, 3, 4] + assert read(parts[0]).column("pfam__count").to_pylist() == [2, 1, 0, 1] assert list(np.frombuffer(payloads["csr:pfam"], " Date: Sun, 6 Sep 2026 01:03:47 +0200 Subject: [PATCH 08/31] fix(utils): rebuild CSR score and evidence payloads with the EAT overlay They are numbered by global hit, so a predicted row shifted every later one. --- .../src/visualization/eat-overlay.test.ts | 77 ++++++++ .../utils/src/visualization/eat-overlay.ts | 179 ++++++++++++++---- 2 files changed, 218 insertions(+), 38 deletions(-) diff --git a/packages/utils/src/visualization/eat-overlay.test.ts b/packages/utils/src/visualization/eat-overlay.test.ts index 22970238..289728aa 100644 --- a/packages/utils/src/visualization/eat-overlay.test.ts +++ b/packages/utils/src/visualization/eat-overlay.test.ts @@ -15,6 +15,7 @@ import { normalizeReliability, parseEatCompanionColumn, } from './eat-overlay'; +import { getProteinEvidence, getProteinScores } from './plot-data-accessors'; function createData(): VisualizationData { return { @@ -234,6 +235,82 @@ describe('EAT overlay over CSR storage (bundle format v3)', () => { expect(out.annotation_data.ec.codes.buffer).not.toBe(source.codes.buffer); }); + // p0 → [A], p1 → [B], p2 → [C], p3 → [A, B], with a score and an evidence code + // attached to every hit. p1's single hit is replaced by a two-valued transfer, which + // renumbers every hit from p2 onwards — the case that silently shifted the payloads. + function createCsrDataWithPayloads(): VisualizationData { + return { + protein_ids: ['p0', 'p1', 'p2', 'p3'], + projections: [{ name: 'umap', dimension: 2, data: new Float32Array(8) }], + annotations: { + ec: { + kind: 'categorical', + values: ['A', 'B', 'C'], + colors: ['#f00', '#0f0', '#00f'], + shapes: ['circle', 'circle', 'circle'], + }, + }, + annotation_data: { + ec: { + kind: 'csr', + end: Int32Array.of(1, 2, 3, 5), + codes: Int32Array.of(0, 1, 2, 0, 1), + length: 4, + }, + }, + // Hit 4 deliberately carries neither a score nor an evidence code. + annotation_scores_csr: { + ec: { + hitEnd: Int32Array.of(1, 2, 3, 4, 4), + values: Float32Array.of(0.25, 0.5, 0.75, 0.125), + }, + }, + annotation_evidence_csr: { + ec: { codes: Int32Array.of(0, 1, 2, 0, -1), dict: ['IDA', 'IEA', 'ISS'] }, + }, + annotation_predicted: { + ec: [null, { value: 'A;C', values: ['A', 'C'], confidence: 0.6, source: 'p0' }, null, null], + }, + }; + } + + it('renumbers the flat score and evidence payloads with the rebuilt hits', () => { + const out = materializeEatOverlay(createCsrDataWithPayloads(), 'ec', true); + const rows = out.annotation_data.ec; + if (!isCsrAnnotationData(rows)) throw new Error('expected CSR'); + expect(Array.from(rows.end)).toEqual([1, 3, 4, 6]); + + // Curated rows keep their own score and evidence, including the rows AFTER the + // multi-valued transfer that shifted every later hit number. + expect(getProteinScores(out, 0, 'ec')).toEqual([[0.25]]); + expect(getProteinEvidence(out, 0, 'ec')).toEqual(['IDA']); + expect(getProteinScores(out, 2, 'ec')).toEqual([[0.75]]); + expect(getProteinEvidence(out, 2, 'ec')).toEqual(['ISS']); + expect(getProteinScores(out, 3, 'ec')).toEqual([[0.125], null]); + expect(getProteinEvidence(out, 3, 'ec')).toEqual(['IDA', null]); + + // The predicted row carries no curated score or evidence of its own, one slot per value. + expect(getProteinScores(out, 1, 'ec')).toEqual([null, null]); + expect(getProteinEvidence(out, 1, 'ec')).toEqual([null, null]); + }); + + it('leaves the curated payloads untouched and drops no unused capacity', () => { + const data = createCsrDataWithPayloads(); + const out = materializeEatOverlay(data, 'ec', true); + expect(out.annotation_scores_csr?.ec).not.toBe(data.annotation_scores_csr?.ec); + expect(Array.from(data.annotation_scores_csr!.ec.values)).toEqual([0.25, 0.5, 0.75, 0.125]); + expect(Array.from(data.annotation_evidence_csr!.ec.codes)).toEqual([0, 1, 2, 0, -1]); + // p1's 0.5 is gone with the row it belonged to; the slack is not retained. + expect(Array.from(out.annotation_scores_csr!.ec.values)).toEqual([0.25, 0.75, 0.125]); + expect(out.annotation_scores_csr!.ec.values.buffer.byteLength).toBe(3 * 4); + }); + + it('does not grow the payload records on a dataset that carries none', () => { + const out = materializeEatOverlay(createCsrData(), 'ec', true); + expect('annotation_scores_csr' in out).toBe(false); + expect('annotation_evidence_csr' in out).toBe(false); + }); + it('leaves the row alone when a prediction resolves to no known value', () => { const data = createCsrData(); data.annotation_predicted!.ec = [ diff --git a/packages/utils/src/visualization/eat-overlay.ts b/packages/utils/src/visualization/eat-overlay.ts index 9d5f338a..f5664493 100644 --- a/packages/utils/src/visualization/eat-overlay.ts +++ b/packages/utils/src/visualization/eat-overlay.ts @@ -1,4 +1,11 @@ -import type { AnnotationData, PredictedCell, VisualizationData } from '../types.js'; +import type { + AnnotationData, + CsrAnnotationData, + CsrEvidence, + CsrScores, + PredictedCell, + VisualizationData, +} from '../types.js'; import { getFirstAnnotationIndex, getProteinAnnotationIndices, @@ -196,8 +203,106 @@ function predictedIndices( .filter((index) => index >= 0); } +/** Which rows a prediction actually replaces, and with which value indices. */ +function predictedReplacements( + predictedCells: readonly (PredictedCell | null)[], + valueToIndex: ReadonlyMap, + rowCount: number, +): Map { + const replacements = new Map(); + for (let i = 0; i < predictedCells.length && i < rowCount; i++) { + const cell = predictedCells[i]; + if (!cell) continue; + const indices = predictedIndices(cell, valueToIndex); + if (indices.length > 0) replacements.set(i, indices); + } + return replacements; +} + +/** + * CSR rebuild for the overlay — codes AND the flat per-hit payloads, in one lockstep pass. + * + * Rebuilt as CSR, never densified to `number[][]`: a 573K column would otherwise become one + * array per protein. Predicted rows replace the source row's hits wholesale; a prediction + * that resolves to nothing leaves the row alone, as the Int32Array and sparse-multi branches + * do (CSR has no way to store the `-1` the dense branch falls back to). + * + * `annotation_scores_csr` / `annotation_evidence_csr` are numbered by GLOBAL hit, so a + * predicted row whose hit count differs from the curated one renumbers every LATER hit. + * Rebuilding the codes alone left those payloads addressing other proteins' hits, and + * silently: the accessors read past the end of `hitEnd` / `codes` and return null instead of + * throwing. So they are rebuilt here alongside the codes — a preserved row copies its hit + * spans verbatim, a replaced row emits one empty score range and one `-1` evidence code per + * predicted value (an EAT transfer carries no curated score or evidence of its own). + */ +function cloneCsrWithPredictions( + source: CsrAnnotationData, + predictedCells: readonly (PredictedCell | null)[], + valueToIndex: ReadonlyMap, + sourceScores: CsrScores | undefined, + sourceEvidence: CsrEvidence | undefined, +): { rows: CsrAnnotationData; scores?: CsrScores; evidence?: CsrEvidence } { + const replacements = predictedReplacements(predictedCells, valueToIndex, source.length); + + const end = new Int32Array(source.length); + let total = 0; + for (let i = 0; i < source.length; i++) { + const replacement = replacements.get(i); + total += replacement ? replacement.length : source.end[i] - (i === 0 ? 0 : source.end[i - 1]); + end[i] = total; + } + + const codes = new Int32Array(total); + // The payloads can only shrink (a replaced row drops its own values), so the source + // length is a safe upper bound and the trailing slack is sliced off at the end. + const values = new Float32Array(sourceScores ? sourceScores.values.length : 0); + const hitEnd = new Int32Array(sourceScores ? total : 0); + const evidenceCodes = new Int32Array(sourceEvidence ? total : 0); + + let dst = 0; + let written = 0; + for (let i = 0; i < source.length; i++) { + const start = i === 0 ? 0 : source.end[i - 1]; + const stop = source.end[i]; + const replacement = replacements.get(i); + if (replacement) { + for (const index of replacement) { + codes[dst] = index; + // Repeating the running cursor is the empty range that says "no score". + if (sourceScores) hitEnd[dst] = written; + if (sourceEvidence) evidenceCodes[dst] = -1; + dst++; + } + continue; + } + if (!sourceScores && !sourceEvidence) { + // No payload to walk, so a preserved row is one bulk copy. + if (stop > start) codes.set(source.codes.subarray(start, stop), dst); + dst += stop - start; + continue; + } + for (let hit = start; hit < stop; hit++) { + codes[dst] = source.codes[hit]; + if (sourceScores) { + const from = hit === 0 ? 0 : sourceScores.hitEnd[hit - 1]; + const to = sourceScores.hitEnd[hit]; + for (let v = from; v < to; v++) values[written++] = sourceScores.values[v]; + hitEnd[dst] = written; + } + if (sourceEvidence) evidenceCodes[dst] = sourceEvidence.codes[hit]; + dst++; + } + } + + return { + rows: { kind: 'csr', end, codes, length: source.length }, + ...(sourceScores ? { scores: { hitEnd, values: values.slice(0, written) } } : {}), + ...(sourceEvidence ? { evidence: { codes: evidenceCodes, dict: sourceEvidence.dict } } : {}), + }; +} + function cloneWithPredictions( - source: AnnotationData, + source: Exclude, predictedCells: readonly (PredictedCell | null)[], valueToIndex: ReadonlyMap, ): AnnotationData { @@ -247,42 +352,6 @@ function cloneWithPredictions( return { kind: 'sparse-multi', base, overrides, length: base.length }; } - if (isCsrAnnotationData(source)) { - // Rebuilt as CSR, never densified to `number[][]`: a 573K column would otherwise - // become one array per protein. Predicted rows replace the source row's hits - // wholesale; a prediction that resolves to nothing leaves the row alone, as the - // Int32Array and sparse-multi branches above do (CSR has no way to store the - // `-1` the dense branch falls back to). - const replacements = new Map(); - for (let i = 0; i < predictedCells.length && i < source.length; i++) { - const cell = predictedCells[i]; - if (!cell) continue; - const indices = predictedIndices(cell, valueToIndex); - if (indices.length > 0) replacements.set(i, indices); - } - const end = new Int32Array(source.length); - let total = 0; - for (let i = 0; i < source.length; i++) { - const replacement = replacements.get(i); - total += replacement ? replacement.length : source.end[i] - (i === 0 ? 0 : source.end[i - 1]); - end[i] = total; - } - const codes = new Int32Array(total); - let cursor = 0; - for (let i = 0; i < source.length; i++) { - const replacement = replacements.get(i); - if (replacement) { - for (const index of replacement) codes[cursor++] = index; - continue; - } - const start = i === 0 ? 0 : source.end[i - 1]; - const stop = source.end[i]; - if (stop > start) codes.set(source.codes.subarray(start, stop), cursor); - cursor += stop - start; - } - return { kind: 'csr', end, codes, length: source.length }; - } - const clone = source.slice(); for (let i = 0; i < predictedCells.length; i++) { const cell = predictedCells[i]; @@ -314,6 +383,40 @@ export function materializeEatOverlay( if (value != null) valueToIndex.set(value, index); }); + if (isCsrAnnotationData(source)) { + // The payloads come back from the same pass as the codes and are returned here rather + // than left to the `...data` spread: a stale `annotation_scores_csr` carried through by + // reference is numbered against the OLD hit layout, which shows one protein's score and + // evidence on the next protein's tooltip without any visible failure. + const rebuilt = cloneCsrWithPredictions( + source, + predictedCells, + valueToIndex, + data.annotation_scores_csr?.[annotationKey], + data.annotation_evidence_csr?.[annotationKey], + ); + return { + ...data, + annotation_data: { ...data.annotation_data, [annotationKey]: rebuilt.rows }, + ...(rebuilt.scores + ? { + annotation_scores_csr: { + ...data.annotation_scores_csr, + [annotationKey]: rebuilt.scores, + }, + } + : {}), + ...(rebuilt.evidence + ? { + annotation_evidence_csr: { + ...data.annotation_evidence_csr, + [annotationKey]: rebuilt.evidence, + }, + } + : {}), + }; + } + return { ...data, annotation_data: { From 41cfb9588eaf0d587dd501c28e01602c04f97114 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 01:04:11 +0200 Subject: [PATCH 09/31] fix(core): throw when a CSR remap drops a hit instead of desyncing payloads --- .../data-loader/utils/conversion.test.ts | 37 +++++++++++++------ .../data-loader/utils/conversion.ts | 17 +++++++-- 2 files changed, 39 insertions(+), 15 deletions(-) diff --git a/packages/core/src/components/data-loader/utils/conversion.test.ts b/packages/core/src/components/data-loader/utils/conversion.test.ts index f5a1a387..72ba05cd 100644 --- a/packages/core/src/components/data-loader/utils/conversion.test.ts +++ b/packages/core/src/components/data-loader/utils/conversion.test.ts @@ -741,11 +741,11 @@ describe('splitCategoricalAnnotationValues v2', () => { describe('normalizeEatCompanionColumns over CSR storage (bundle format v3)', () => { /** - * Base column `ec` in CSR, values `['A', null, 'B']` — the `null` slot is what makes - * the remap drop a code, so the rebuilt payload has to compact: + * Base column `ec` in CSR, values `['A', null, 'B']`. Nothing points at the `null` + * slot, so every hit survives the remap and only the value INDICES shift (B: 2 → 1): * p0 → [] (curated missing → the EAT companion applies) * p1 → [A] - * p2 → [null, B] + * p2 → [B] * p3 → [A, B] * The companion `ec__pred_value` is CSR too, with its scores/evidence in the flat * v3 payloads rather than the nested records. @@ -778,8 +778,8 @@ describe('normalizeEatCompanionColumns over CSR storage (bundle format v3)', () annotation_data: { ec: { kind: 'csr', - end: Int32Array.of(0, 1, 3, 5), - codes: Int32Array.of(0, 1, 2, 0, 2), + end: Int32Array.of(0, 1, 2, 4), + codes: Int32Array.of(0, 2, 0, 2), length: 4, }, ec__pred_value: { @@ -805,29 +805,42 @@ describe('normalizeEatCompanionColumns over CSR storage (bundle format v3)', () }; } - it('remaps CSR storage into fresh compacted buffers', () => { + it('remaps CSR storage into fresh buffers', () => { const src = csrEatData(); const sourceRows = src.annotation_data.ec as CsrAnnotationData; const out = normalizeEatCompanionColumns(src); const rows = out.annotation_data.ec; if (!isCsrAnnotationData(rows)) throw new Error('expected CSR storage'); - // 'A','B' survive, the null slot is dropped, 'C' is appended by the transfer. + // 'A','B' survive, the unreferenced null slot leaves the value list, 'C' is + // appended by the transfer — so 'B' moves from index 2 to index 1. expect(out.annotations.ec.values).toEqual(['A', 'B', 'C']); // The transfer itself stays in `annotation_predicted`; only the curated rows // are remapped here (materializeEatOverlay applies the prediction later). expect(getProteinAnnotationIndices(rows, 0)).toEqual([]); expect(getProteinAnnotationIndices(rows, 1)).toEqual([0]); // 'A' - expect(getProteinAnnotationIndices(rows, 2)).toEqual([1]); // null dropped, 'B' kept + expect(getProteinAnnotationIndices(rows, 2)).toEqual([1]); // 'B', renumbered expect(getProteinAnnotationIndices(rows, 3)).toEqual([0, 1]); expect(Array.from(rows.end)).toEqual([0, 1, 2, 4]); - // Compacted: one code fewer than the source, and the trailing slack is not - // retained — `.slice(0, written)` hands back an exactly-sized fresh buffer. - expect(rows.codes.length).toBe(4); + // Neither the worst-case buffer nor the source's is retained — `.slice(0, written)` + // hands back an exactly-sized fresh one. expect(rows.codes.buffer.byteLength).toBe(4 * 4); expect(rows.codes.buffer).not.toBe(sourceRows.codes.buffer); - expect(Array.from(sourceRows.codes)).toEqual([0, 1, 2, 0, 2]); + expect(Array.from(sourceRows.codes)).toEqual([0, 2, 0, 2]); + }); + + it('throws rather than silently renumbering the flat payloads when a hit drops', () => { + const src = csrEatData(); + // p1's hit now points at the null value slot, which the remap drops. Every later + // hit number would shift under the unchanged annotation_scores_csr / _evidence_csr. + src.annotation_data.ec = { + kind: 'csr', + end: Int32Array.of(0, 1, 2, 4), + codes: Int32Array.of(1, 2, 0, 2), + length: 4, + }; + expect(() => normalizeEatCompanionColumns(src)).toThrow(/CSR remap dropped 1 of 4 hits/); }); it('reads the companion column scores/evidence from the flat v3 payloads', () => { diff --git a/packages/core/src/components/data-loader/utils/conversion.ts b/packages/core/src/components/data-loader/utils/conversion.ts index f63cbd22..5cc4588c 100644 --- a/packages/core/src/components/data-loader/utils/conversion.ts +++ b/packages/core/src/components/data-loader/utils/conversion.ts @@ -226,9 +226,8 @@ function remapCategoricalStorage( return { kind: 'sparse-multi', base, overrides, length: base.length }; } if (isCsrAnnotationData(source)) { - // Remapping can drop a hit (`remap` → -1), so the row boundaries shift and the - // payload shrinks: both arrays are rebuilt from scratch, and `codes` is trimmed - // to what was written so the worst-case buffer is neither retained nor shared. + // Both arrays are rebuilt from scratch, and `codes` is trimmed to what was written so + // the worst-case buffer is neither retained nor shared. const end = new Int32Array(source.length); const codes = new Int32Array(source.codes.length); let written = 0; @@ -240,6 +239,18 @@ function remapCategoricalStorage( } end[i] = written; } + if (written !== source.codes.length) { + // Dropping a hit renumbers every later one, and `annotation_scores_csr` / + // `annotation_evidence_csr` are numbered by GLOBAL hit — so a silent drop would + // show one protein's score and evidence on another's tooltip. No shipping producer + // can reach this (`values` is always built as `string[]`), so failing loudly beats + // implementing a renumbering pass nothing exercises. + throw new Error( + `CSR remap dropped ${source.codes.length - written} of ${source.codes.length} hits ` + + '(a null or unmapped annotation value). Renumbering the parallel ' + + 'annotation_scores_csr / annotation_evidence_csr payloads is not implemented.', + ); + } return { kind: 'csr', end, codes: codes.slice(0, written), length: source.length }; } return source.map((indices) => indices.map(remap).filter((index) => index >= 0)); From 86fe5df10cc84ce3b2b66dada9e2cc9981eabd5e Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 01:04:27 +0200 Subject: [PATCH 10/31] fix(web): record an import as pending before the render, copy bytes alongside it Restores the crash-recovery window without putting the byte copy before first paint. --- .../explore/dataset-controller.eat.test.ts | 3 +- .../dataset-controller.persistence.test.ts | 65 ++++++++++++++----- apps/web/src/explore/dataset-controller.ts | 45 +++++++++---- .../src/explore/opfs-dataset-store.test.ts | 38 ++++++++++- apps/web/src/explore/opfs-dataset-store.ts | 36 +++++++++- 5 files changed, 155 insertions(+), 32 deletions(-) diff --git a/apps/web/src/explore/dataset-controller.eat.test.ts b/apps/web/src/explore/dataset-controller.eat.test.ts index 55efe49a..85d7bd11 100644 --- a/apps/web/src/explore/dataset-controller.eat.test.ts +++ b/apps/web/src/explore/dataset-controller.eat.test.ts @@ -24,7 +24,8 @@ vi.mock('./persisted-dataset', () => ({ vi.mock('./opfs-dataset-store', () => ({ markLastLoadStatus: mocks.markLastLoadStatus, - saveLastImportedFile: vi.fn(), + saveLastImportedFileMetadata: vi.fn(), + saveLastImportedFileData: vi.fn(), })); vi.mock('./tooltip-annotations-store', () => ({ diff --git a/apps/web/src/explore/dataset-controller.persistence.test.ts b/apps/web/src/explore/dataset-controller.persistence.test.ts index 15841b51..4a774a76 100644 --- a/apps/web/src/explore/dataset-controller.persistence.test.ts +++ b/apps/web/src/explore/dataset-controller.persistence.test.ts @@ -5,7 +5,8 @@ import { createEmptyExploreViewRequest } from './url-state'; const mocks = vi.hoisted(() => ({ loadData: vi.fn(), markLastLoadStatus: vi.fn(), - saveLastImportedFile: vi.fn(), + saveLastImportedFileMetadata: vi.fn(), + saveLastImportedFileData: vi.fn(), resolvePendingLoadFinalization: vi.fn(), warning: vi.fn(), })); @@ -26,7 +27,8 @@ vi.mock('./persisted-dataset', () => ({ vi.mock('./opfs-dataset-store', () => ({ markLastLoadStatus: mocks.markLastLoadStatus, - saveLastImportedFile: mocks.saveLastImportedFile, + saveLastImportedFileMetadata: mocks.saveLastImportedFileMetadata, + saveLastImportedFileData: mocks.saveLastImportedFileData, })); vi.mock('./tooltip-annotations-store', () => ({ @@ -96,44 +98,61 @@ describe('dataset controller OPFS persistence', () => { beforeEach(() => { vi.clearAllMocks(); mocks.markLastLoadStatus.mockResolvedValue(undefined); - mocks.saveLastImportedFile.mockResolvedValue(undefined); + mocks.saveLastImportedFileMetadata.mockResolvedValue(undefined); + mocks.saveLastImportedFileData.mockResolvedValue(undefined); }); - it('saves the imported file only after the render pass resolves', async () => { + /** Drain the microtask queue so every already-resolved await has run. */ + const flush = () => new Promise((resolve) => setTimeout(resolve, 0)); + + it('opens the pending window before the render and settles the byte copy after it', async () => { let finishRender = () => {}; + let finishCopy = () => {}; mocks.loadData.mockImplementation( () => new Promise((resolve) => { finishRender = () => resolve(); }), ); + mocks.saveLastImportedFileData.mockImplementation( + () => + new Promise((resolve) => { + finishCopy = () => resolve(); + }), + ); const { controller, overlayController } = buildController(); const pending = controller.handleDataLoaded(loadedEvent); + await flush(); - await Promise.resolve(); - await Promise.resolve(); + // The crash-recovery window is open before the render it has to survive: a tab that + // dies here leaves a `pending` record naming this import, not the previous dataset's. + expect(mocks.saveLastImportedFileMetadata).toHaveBeenCalledWith(file); + expect(mocks.saveLastImportedFileMetadata.mock.invocationCallOrder[0]).toBeLessThan( + mocks.loadData.mock.invocationCallOrder[0], + ); + // The byte copy is in flight rather than in front of first paint: the render started + // even though the copy has not resolved, and there is no blocking overlay. + expect(mocks.saveLastImportedFileData).toHaveBeenCalledWith(file); expect(mocks.loadData).toHaveBeenCalledOnce(); - expect(mocks.saveLastImportedFile).not.toHaveBeenCalled(); - // No blocking "Saving imported dataset..." overlay in front of the render. expect(overlayController.update).not.toHaveBeenCalled(); + // ...and success is not reported over a half-copied file. finishRender(); + await flush(); + expect(mocks.markLastLoadStatus).not.toHaveBeenCalled(); + + finishCopy(); await pending; - expect(mocks.saveLastImportedFile).toHaveBeenCalledWith(file); - // markLastLoadStatus reads the metadata the save writes, so it must run after. - expect(mocks.saveLastImportedFile.mock.invocationCallOrder[0]).toBeLessThan( - mocks.markLastLoadStatus.mock.invocationCallOrder[0], - ); expect(mocks.markLastLoadStatus).toHaveBeenCalledWith('success'); expect(mocks.resolvePendingLoadFinalization).toHaveBeenCalledWith(3); }); - it('warns but still finishes the load when the save fails', async () => { + it('warns but still finishes the load when the byte copy fails', async () => { const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}); mocks.loadData.mockResolvedValue(undefined); - mocks.saveLastImportedFile.mockRejectedValue(new Error('quota exceeded')); + mocks.saveLastImportedFileData.mockRejectedValue(new Error('quota exceeded')); const { controller } = buildController(); await controller.handleDataLoaded(loadedEvent); @@ -143,4 +162,20 @@ describe('dataset controller OPFS persistence', () => { expect(mocks.resolvePendingLoadFinalization).toHaveBeenCalledWith(3); consoleError.mockRestore(); }); + + it('warns and still renders when the pending record cannot be written', async () => { + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}); + mocks.loadData.mockResolvedValue(undefined); + mocks.saveLastImportedFileMetadata.mockRejectedValue(new Error('quota exceeded')); + + const { controller } = buildController(); + await controller.handleDataLoaded(loadedEvent); + + // No recovery window, but the dataset still reaches the screen. + expect(mocks.saveLastImportedFileData).not.toHaveBeenCalled(); + expect(mocks.loadData).toHaveBeenCalledOnce(); + expect(mocks.warning).toHaveBeenCalledOnce(); + expect(mocks.resolvePendingLoadFinalization).toHaveBeenCalledWith(3); + consoleError.mockRestore(); + }); }); diff --git a/apps/web/src/explore/dataset-controller.ts b/apps/web/src/explore/dataset-controller.ts index 9b2721dd..099f7701 100644 --- a/apps/web/src/explore/dataset-controller.ts +++ b/apps/web/src/explore/dataset-controller.ts @@ -13,7 +13,11 @@ import { getDataLoadFailureNotification, getDatasetPersistenceFailureNotification, } from './notifications'; -import { markLastLoadStatus, saveLastImportedFile } from './opfs-dataset-store'; +import { + markLastLoadStatus, + saveLastImportedFileData, + saveLastImportedFileMetadata, +} from './opfs-dataset-store'; import { createDataRenderer } from './data-renderer'; import type { InteractionController } from './interaction-controller'; import type { LoadQueue } from './load-queue'; @@ -123,6 +127,26 @@ export function createDatasetController({ legendElement.clearForNewDataset(datasetHash, shouldClearPersistedState); controlBar.clearForNewDataset(datasetHash, shouldClearPersistedState); + // The `pending` record is written and awaited BEFORE the render: that window is what + // the recovery banner reads, so a tab that dies mid-render leaves a record of the + // import that died instead of silently restoring the previous dataset. It is also + // what markLastLoadStatus reads on both the success and the data-error path. + // + // The byte copy is a different matter — 45-145 MB, and it used to sit in front of + // first paint. It is only STARTED here; the await happens after the render below. + let persistBytes: Promise | null = null; + if (loadMeta.kind === 'user' && file) { + try { + await saveLastImportedFileMetadata(file); + persistBytes = saveLastImportedFileData(file).then( + () => null, + (error: unknown) => error, + ); + } catch (error) { + persistBytes = Promise.resolve(error); + } + } + await loadData(data); if (settings && loadMeta.kind !== 'opfs') { @@ -221,18 +245,13 @@ export function createDatasetController({ viewController.applyLatestViewForDatasetLoad(data); - // Persisted only once the dataset is on screen: this is a full byte copy of - // the bundle (45-145 MB) into OPFS and used to sit in front of first paint. - // Load-queue serialisation (it waits on resolvePendingLoadFinalization below) - // keeps a later import from interleaving with this write, and - // markLastLoadStatus reads the metadata written here, so it must stay after. - if (loadMeta.kind === 'user' && file) { - try { - await saveLastImportedFile(file); - } catch (error) { - console.error('Failed to persist imported dataset in OPFS:', error); - notify.warning(getDatasetPersistenceFailureNotification(error)); - } + // Settled only once the dataset is on screen. Load-queue serialisation (it waits on + // resolvePendingLoadFinalization below) keeps a later import from interleaving with + // this write, and markLastLoadStatus must not report success over a half-copied file. + const persistError = await persistBytes; + if (persistError) { + console.error('Failed to persist imported dataset in OPFS:', persistError); + notify.warning(getDatasetPersistenceFailureNotification(persistError)); } try { diff --git a/apps/web/src/explore/opfs-dataset-store.test.ts b/apps/web/src/explore/opfs-dataset-store.test.ts index 35604376..39208547 100644 --- a/apps/web/src/explore/opfs-dataset-store.test.ts +++ b/apps/web/src/explore/opfs-dataset-store.test.ts @@ -6,9 +6,16 @@ import { loadLastImportedFile, markLastLoadStatus, readLastLoadStatus, - saveLastImportedFile, + saveLastImportedFileData, + saveLastImportedFileMetadata, } from './opfs-dataset-store'; +/** Both halves back to back — what a completed import leaves on disk. */ +async function saveLastImportedFile(file: File): Promise { + await saveLastImportedFileMetadata(file); + await saveLastImportedFileData(file); +} + class MockWritableFileStream { private chunks: BlobPart[] = []; private onClose: (blob: Blob) => void; @@ -141,6 +148,35 @@ describe('opfs-dataset-store', () => { expect(await loaded?.text()).toBe('protein-data'); }); + it('records the pending import before any bytes are copied', async () => { + const root = new MockDirectoryHandle(); + stubNavigator(root); + + await saveLastImportedFileMetadata(new File(['x'], 'first.parquetbundle')); + + // The crash-recovery window is open on the metadata alone: this is the state a tab + // that dies during the render leaves behind. + expect(await readLastLoadStatus()).toEqual({ + status: 'pending', + lastError: undefined, + failedAttempts: 0, + }); + }); + + it('drops the previous bytes with the previous metadata', async () => { + const root = new MockDirectoryHandle(); + stubNavigator(root); + + await saveLastImportedFile(new File(['first-data'], 'first.parquetbundle')); + await saveLastImportedFileMetadata(new File(['second-data'], 'second.parquetbundle')); + + // Without the drop the first dataset's bytes would come back named `second`, which is + // a different dataset restored silently. (Reading in this state also clears the store, + // so a crash mid-copy costs the previous dataset — as it always did, since the copy + // overwrote it in place.) + expect(await loadLastImportedFile()).toBeNull(); + }); + it('returns null and clears the store when the payload file is missing', async () => { const root = new MockDirectoryHandle(); const store = await root.getDirectoryHandle('protspace-last-import', { create: true }); diff --git a/apps/web/src/explore/opfs-dataset-store.ts b/apps/web/src/explore/opfs-dataset-store.ts index 746e961c..84669c32 100644 --- a/apps/web/src/explore/opfs-dataset-store.ts +++ b/apps/web/src/explore/opfs-dataset-store.ts @@ -176,7 +176,7 @@ export function isSupported(): boolean { return typeof navigator !== 'undefined' && hasStorageDirectoryApi(navigator.storage); } -export async function saveLastImportedFile(file: File): Promise { +async function requireStoreDirectory(): Promise { if (!isSupported()) { throw buildSupportError(); } @@ -185,6 +185,22 @@ export async function saveLastImportedFile(file: File): Promise { if (!directory) { throw new Error('Unable to access the Origin Private File System.'); } + return directory; +} + +/** + * Open the crash-recovery window for an import: the `pending` record only, no bytes. + * + * Split from the byte copy below so the import path can record the file BEFORE the render. + * A tab that dies mid-render is exactly what `pending` and the recovery banner exist for, + * and the copy is far too large (45-145 MB) to sit in front of first paint. + * + * The previous dataset's bytes go with the previous metadata. Left in place they would be + * handed back under THIS file's name if the copy never finished — a different dataset, + * silently, which is worse than restoring nothing. + */ +export async function saveLastImportedFileMetadata(file: File): Promise { + const directory = await requireStoreDirectory(); const metadata: StoredDatasetMetadata = { schemaVersion: SCHEMA_VERSION, @@ -198,8 +214,24 @@ export async function saveLastImportedFile(file: File): Promise { }; try { - await writeBlobFile(directory, DATA_FILENAME, file); await writeMetadata(directory, metadata); + try { + await directory.removeEntry(DATA_FILENAME); + } catch { + // Nothing stored yet. + } + } catch (error) { + await clearStoreDirectory(); + throw error instanceof Error ? error : new Error('Failed to save imported dataset.'); + } +} + +/** Byte-copy half of the import save. Runs after {@link saveLastImportedFileMetadata}. */ +export async function saveLastImportedFileData(file: File): Promise { + const directory = await requireStoreDirectory(); + + try { + await writeBlobFile(directory, DATA_FILENAME, file); } catch (error) { await clearStoreDirectory(); throw error instanceof Error ? error : new Error('Failed to save imported dataset.'); From 4a11f2590ce06a28a9ca5071806fe60c4a1632ea Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 01:04:43 +0200 Subject: [PATCH 11/31] perf(utils): skip the per-protein score and evidence reads on columns with neither --- packages/utils/src/parquet/bundle-writer.ts | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/packages/utils/src/parquet/bundle-writer.ts b/packages/utils/src/parquet/bundle-writer.ts index 796f005f..f0e57ca0 100644 --- a/packages/utils/src/parquet/bundle-writer.ts +++ b/packages/utils/src/parquet/bundle-writer.ts @@ -33,6 +33,10 @@ import { getEatCompanionColumn, getPredictedCellValues } from '../visualization/ import { isNAValue } from '../visualization/missing-values.js'; import { encodeAnnotationField } from './annotation-codec.js'; +/** Stand-in for a column with no scores or evidence: every index reads `undefined`, exactly + * as the empty array the accessors would have allocated per protein. */ +const NO_HITS: readonly never[] = []; + const ANNOTATION_FORMAT_VERSION = '2'; const ANNOTATION_FORMAT_VERSION_KEY = 'protspace_format_version'; @@ -134,6 +138,16 @@ function createAnnotationsParquet(data: VisualizationData): ArrayBuffer { const annotationIndices = data.annotation_data[annotationName]; if (!annotationIndices) continue; + // Hoisted out of the protein loop: both accessors allocate a fresh array on every + // call, including the v1/v2 case where the dataset carries no scores or evidence at + // all — 573K proteins x ~24 categorical columns of throwaway arrays per export. + const hasScores = Boolean( + data.annotation_scores?.[annotationName] ?? data.annotation_scores_csr?.[annotationName], + ); + const hasEvidence = Boolean( + data.annotation_evidence?.[annotationName] ?? data.annotation_evidence_csr?.[annotationName], + ); + // Convert indices back to actual annotation values const values: (string | null)[] = new Array(data.protein_ids.length); for (let i = 0; i < data.protein_ids.length; i++) { @@ -149,8 +163,8 @@ function createAnnotationsParquet(data: VisualizationData): ArrayBuffer { // records: a v3 (CSR) dataset carries these flat per hit, and only the accessors // know both layouts. Both return per-hit arrays in `getProteinAnnotationIndices` // order, so `cellIndex` still lines up. - const proteinEvidence = getProteinEvidence(data, i, annotationName); - const proteinScores = getProteinScores(data, i, annotationName); + const proteinEvidence = hasEvidence ? getProteinEvidence(data, i, annotationName) : NO_HITS; + const proteinScores = hasScores ? getProteinScores(data, i, annotationName) : NO_HITS; const cellValues = getProteinAnnotationIndices(annotationIndices, i).flatMap( (valueIndex, cellIndex) => { const value = annotation.values[valueIndex]; From 85239fdbe76ee54c822f6865f5e1667903f26ce6 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 01:04:54 +0200 Subject: [PATCH 12/31] docs(core): document the grid-vs-d3 divergences and cover the NaN depth fallback --- .../interaction/quadtree-index.ts | 9 +++++++- .../webgl/renderer/depth-sort.test.ts | 22 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/packages/core/src/components/scatter-plot/interaction/quadtree-index.ts b/packages/core/src/components/scatter-plot/interaction/quadtree-index.ts index 50fc0066..a7c87e60 100644 --- a/packages/core/src/components/scatter-plot/interaction/quadtree-index.ts +++ b/packages/core/src/components/scatter-plot/interaction/quadtree-index.ts @@ -8,13 +8,20 @@ import type { PlotData } from '@protspace/utils'; * allocated one heap object per visible point, which cost 512 ms to build at 573K points * against 37 ms for the grid. The class name is kept because every consumer refers to it. * - * Behaviour is intentionally identical to the quadtree it replaced: + * Behaviour matches the quadtree it replaced: * - non-finite screen coordinates are not indexed; * - `queryByPixels` / `queryByPolygon` use an inclusive AABB test and return every match, * including coincident points (result order is not contractual); * - `findNearest` uses a strict `< radius` cutoff and, for exactly coincident points, returns * the LAST one in `slots` order, which is what `d3.quadtree.find` did (a coincident point * is pushed to the head of the leaf chain by `add`, and `find` only reads the head). + * + * Two cases diverge, both harmless and both left as they are: + * - among equidistant but NON-coincident points d3 returns the first it visits and this + * returns the last (`d2 <= best` below). Which point of several at the same distance wins + * a hover is arbitrary either way; + * - a negative radius returns a hit here where d3 returned -1. Unreachable: every caller + * passes a positive pixel radius. */ /** Upper bound on grid cells per axis, so a pathological screen extent cannot blow up memory. */ diff --git a/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.test.ts b/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.test.ts index b7feacca..7f86f666 100644 --- a/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.test.ts +++ b/packages/core/src/components/scatter-plot/webgl/renderer/depth-sort.test.ts @@ -119,6 +119,28 @@ describe('sortIndicesByDepthDescending parity with the comparator', () => { expect(Array.from(order)).toEqual(referenceOrder(depths, n)); }); + it('falls back to the comparator on a NaN depth', () => { + // NaN makes `depths[b] - depths[a]` NaN, so the comparator falls through to the index + // tiebreak and index order wins. The counting sort would instead give NaN a bucket of + // its own — `Array.from(rank.keys()).sort()` leaves it after 1 — and return [0, 2, 1]. + const depths = new Float32Array([1, Number.NaN, 1]); + const order = new Uint32Array(3); + sortIndicesByDepthDescending(order, depths, 3); + expect(Array.from(order)).toEqual([0, 1, 2]); + expect(Array.from(order)).toEqual(referenceOrder(depths, 3)); + }); + + it('sorts infinities through the counting-sort path', () => { + const depths = new Float32Array([Infinity, 0, Infinity, -Infinity, 0]); + const order = new Uint32Array(5); + sortIndicesByDepthDescending(order, depths, 5); + expect(Array.from(order)).toEqual([0, 2, 1, 4, 3]); + expect(Array.from(order)).toEqual(referenceOrder(depths, 5)); + }); + + // The cap is not unreachable: `getDepth` yields roughly one distinct value per legend + // slot per opacity, so a high-cardinality categorical column (thousands of categories) + // exceeds 4096 and takes this path in production, not just in tests. it('matches the comparator above the distinct-value cap (fallback path)', () => { const rng = makeRng(777); const n = 20000; From 2759fea8a3addcf87c25f832d4f7ebbd981f28b4 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 01:28:08 +0200 Subject: [PATCH 13/31] perf(core): bincount legend counts from storage instead of a flat array Counting per protein index also fixes isolation-mode count misalignment. --- .../legend/annotation-values.test.ts | 139 ------------ .../components/legend/annotation-values.ts | 34 --- .../legend/legend-data-processor.test.ts | 204 +++++++++++++++++- .../legend/legend-data-processor.ts | 121 +++++++++-- packages/core/src/components/legend/legend.ts | 70 +++++- .../scatter-plot.filter-render.test.ts | 11 +- 6 files changed, 378 insertions(+), 201 deletions(-) delete mode 100644 packages/core/src/components/legend/annotation-values.test.ts delete mode 100644 packages/core/src/components/legend/annotation-values.ts diff --git a/packages/core/src/components/legend/annotation-values.test.ts b/packages/core/src/components/legend/annotation-values.test.ts deleted file mode 100644 index 34840786..00000000 --- a/packages/core/src/components/legend/annotation-values.test.ts +++ /dev/null @@ -1,139 +0,0 @@ -import { describe, it, expect } from 'vitest'; -import { buildAnnotationValueList } from './annotation-values'; -import { NA_VALUE } from './config'; - -// ─── Int32Array (single-valued) fixtures ──────────────────────────────────── - -/** - * Build an Int32Array where each element is the annotation index for that protein. - * -1 means "no annotation". - */ -function makeSingleValued(indices: number[]): Int32Array { - return new Int32Array(indices); -} - -// ─── Multi-valued fixtures ─────────────────────────────────────────────────── - -/** - * Build a multi-valued AnnotationData (readonly (readonly number[])[]) directly. - */ -function makeMultiValued(indexLists: number[][]): readonly (readonly number[])[] { - return indexLists.map((list) => list as readonly number[]); -} - -// ─── Tests ─────────────────────────────────────────────────────────────────── - -describe('buildAnnotationValueList', () => { - describe('empty (proteinCount 0)', () => { - it('returns [] for single-valued storage', () => { - const colData = makeSingleValued([]); - expect(buildAnnotationValueList(colData, [], 0)).toEqual([]); - }); - - it('returns [] for multi-valued storage', () => { - const colData = makeMultiValued([]); - expect(buildAnnotationValueList(colData, [], 0)).toEqual([]); - }); - }); - - describe('single-valued (Int32Array) — all proteins annotated', () => { - it('produces one value per protein in order', () => { - const values = ['alpha', 'beta', 'gamma']; - // protein 0 → values[2], protein 1 → values[0], protein 2 → values[1] - const colData = makeSingleValued([2, 0, 1]); - const result = buildAnnotationValueList(colData, values, 3); - expect(result).toEqual(['gamma', 'alpha', 'beta']); - }); - - it('length equals proteinCount when all are annotated', () => { - const values = ['a', 'b']; - const colData = makeSingleValued([0, 1, 0, 1]); - const result = buildAnnotationValueList(colData, values, 4); - expect(result).toHaveLength(4); - }); - }); - - describe('single-valued (Int32Array) — some proteins missing (compacted)', () => { - it('omits proteins with idx < 0', () => { - const values = ['cat', 'dog']; - // protein 0 → values[0], protein 1 → missing (-1), protein 2 → values[1] - const colData = makeSingleValued([0, -1, 1]); - const result = buildAnnotationValueList(colData, values, 3); - expect(result).toEqual(['cat', 'dog']); - }); - - it('length is less than proteinCount when some are missing', () => { - const values = ['x']; - const colData = makeSingleValued([-1, -1, 0, -1]); - const result = buildAnnotationValueList(colData, values, 4); - expect(result).toHaveLength(1); - expect(result).toEqual(['x']); - }); - - it('preserves order of annotated proteins', () => { - const values = ['first', 'second', 'third']; - // proteins 0,2,4 are annotated; 1,3 are missing - const colData = makeSingleValued([2, -1, 0, -1, 1]); - const result = buildAnnotationValueList(colData, values, 5); - expect(result).toEqual(['third', 'first', 'second']); - }); - - it('returns [] when all proteins have idx < 0', () => { - const values = ['a', 'b']; - const colData = makeSingleValued([-1, -1, -1]); - const result = buildAnnotationValueList(colData, values, 3); - expect(result).toEqual([]); - }); - }); - - describe('multi-valued — expanded', () => { - it('a protein with 2 labels contributes 2 entries', () => { - const values = ['labelA', 'labelB', 'labelC']; - // protein 0 → [0, 1], protein 1 → [2] - const colData = makeMultiValued([[0, 1], [2]]); - const result = buildAnnotationValueList(colData, values, 2); - expect(result).toEqual(['labelA', 'labelB', 'labelC']); - }); - - it('total length equals sum of label counts', () => { - const values = ['v0', 'v1', 'v2', 'v3']; - // protein 0 → 3 labels, protein 1 → 0 labels, protein 2 → 1 label - const colData = makeMultiValued([[0, 1, 2], [], [3]]); - const result = buildAnnotationValueList(colData, values, 3); - expect(result).toHaveLength(4); - expect(result).toEqual(['v0', 'v1', 'v2', 'v3']); - }); - - it('proteins with empty index lists contribute 0 entries', () => { - const values = ['only']; - const colData = makeMultiValued([[], [], [0]]); - const result = buildAnnotationValueList(colData, values, 3); - expect(result).toEqual(['only']); - }); - }); - - describe('toInternalValue applied (null → NA_VALUE)', () => { - it('maps null values to NA_VALUE in single-valued storage', () => { - // values[0] is null → should become NA_VALUE - const values: (string | null)[] = [null, 'real']; - const colData = makeSingleValued([0, 1]); - const result = buildAnnotationValueList(colData, values, 2); - expect(result[0]).toBe(NA_VALUE); - expect(result[1]).toBe('real'); - }); - - it('maps null values to NA_VALUE in multi-valued storage', () => { - const values: (string | null)[] = [null, 'present']; - const colData = makeMultiValued([[0, 1]]); - const result = buildAnnotationValueList(colData, values, 1); - expect(result).toEqual([NA_VALUE, 'present']); - }); - - it('non-null string values pass through unchanged', () => { - const values: (string | null)[] = ['hello', 'world']; - const colData = makeSingleValued([0, 1]); - const result = buildAnnotationValueList(colData, values, 2); - expect(result).toEqual(['hello', 'world']); - }); - }); -}); diff --git a/packages/core/src/components/legend/annotation-values.ts b/packages/core/src/components/legend/annotation-values.ts deleted file mode 100644 index 8da7dae8..00000000 --- a/packages/core/src/components/legend/annotation-values.ts +++ /dev/null @@ -1,34 +0,0 @@ -import type { AnnotationData } from '@protspace/utils'; -import { getFirstAnnotationIndex, getProteinAnnotationIndices } from '@protspace/utils'; -import { toInternalValue } from './config'; - -/** - * Build the flat list of internal annotation values used by the legend frequency count. - * - * Replaces a `protein_ids.flatMap(...)` that allocated a throwaway `[]`/`[value]` array per - * protein. Output is IDENTICAL to that flatMap: - * - Single-valued (Int32Array): one entry per protein that HAS an annotation (idx >= 0), - * in protein order; proteins with no annotation are omitted (compacted). - * - Multi-valued: one entry per (protein, label) pair, in order (expanded). - */ -export function buildAnnotationValueList( - colData: AnnotationData, - values: (string | null)[], - proteinCount: number, -): string[] { - const out: string[] = []; - if (colData instanceof Int32Array) { - for (let i = 0; i < proteinCount; i++) { - const idx = getFirstAnnotationIndex(colData, i); - if (idx >= 0) out.push(toInternalValue(values[idx])); - } - } else { - for (let i = 0; i < proteinCount; i++) { - const indices = getProteinAnnotationIndices(colData, i); - for (let j = 0; j < indices.length; j++) { - out.push(toInternalValue(values[indices[j]])); - } - } - } - return out; -} diff --git a/packages/core/src/components/legend/legend-data-processor.test.ts b/packages/core/src/components/legend/legend-data-processor.test.ts index 5223e96a..be1b4e33 100644 --- a/packages/core/src/components/legend/legend-data-processor.test.ts +++ b/packages/core/src/components/legend/legend-data-processor.test.ts @@ -1,4 +1,6 @@ import { describe, it, expect, beforeEach } from 'vitest'; +import type { AnnotationData } from '@protspace/utils'; +import { getFirstAnnotationIndex, getProteinAnnotationIndices } from '@protspace/utils'; import { LegendDataProcessor, createProcessorContext, @@ -6,7 +8,42 @@ import { } from './legend-data-processor'; import { getVisualEncoding } from './visual-encoding'; import type { LegendItem } from './types'; -import { NA_VALUE, NA_DEFAULT_COLOR } from './config'; +import { NA_VALUE, NA_DEFAULT_COLOR, toInternalValue } from './config'; + +/** + * The pre-bincount path, kept here as the equality reference: materialise a flat + * `string[]` of one interned label per hit, then reduce it to a frequency map. + */ +function legacyCounts( + colData: AnnotationData, + values: (string | null)[], + proteinCount: number, + filteredIndices: Set | null, + knownValues: string[] = [], +): ReadonlyMap { + const list: string[] = []; + if (colData instanceof Int32Array) { + for (let i = 0; i < proteinCount; i++) { + const idx = getFirstAnnotationIndex(colData, i); + if (idx >= 0) list.push(toInternalValue(values[idx])); + } + } else { + for (let i = 0; i < proteinCount; i++) { + for (const idx of getProteinAnnotationIndices(colData, i)) { + list.push(toInternalValue(values[idx])); + } + } + } + return LegendDataProcessor.countAnnotationFrequencies( + list, + filteredIndices !== null, + filteredIndices !== null ? [['isolated']] : [], + filteredIndices ?? new Set(), + knownValues, + ); +} + +const asObject = (counts: ReadonlyMap) => Object.fromEntries(counts); describe('legend-data-processor', () => { let ctx: LegendProcessorContext; @@ -66,6 +103,171 @@ describe('legend-data-processor', () => { }); }); + describe('countFromStorage', () => { + // 6 proteins, 3 declared values, one of which is itself missing ('__NA__'). + // Every fixture below encodes the same shape in a different storage kind: + // p0 -> A, p1 -> B + a null value, p2 -> nothing, p3 -> an out-of-range code, + // p4 -> A + B, p5 -> a negative code inside a hit list. + const VALUES: (string | null)[] = ['A', 'B', null]; + const PROTEINS = 6; + + const dense: readonly (readonly number[])[] = [[0], [1, 2], [], [7], [0, 1], [-1]]; + const int32 = new Int32Array([0, 1, -1, 2, 0, 7]); + const sparse = { + kind: 'sparse-multi' as const, + base: new Int32Array([0, -1, 2, 7, 0, -1]), + overrides: new Map([ + [1, [1, 2]], + [2, []], // an empty override wins over the base code it shadows + [4, [0, 1]], + [5, [-1]], + ]), + length: PROTEINS, + }; + const csr = { + kind: 'csr' as const, + end: new Int32Array([1, 3, 3, 4, 6, 7]), + codes: new Int32Array([0, 1, 2, 7, 0, 1, -1]), + length: PROTEINS, + }; + + // One hit per protein, so the legacy flat array is protein-aligned and its + // isolation filtering is meaningful (see the misalignment cases below). + const alignedInt32 = new Int32Array([0, 1, 2, 0, 1, 2]); + const alignedDense: readonly (readonly number[])[] = [[0], [1], [2], [0], [1], [2]]; + const alignedSparse = { + kind: 'sparse-multi' as const, + base: new Int32Array([0, 1, -1, 0, 1, 2]), + overrides: new Map([[2, [2]]]), + length: PROTEINS, + }; + const alignedCsr = { + kind: 'csr' as const, + end: new Int32Array([1, 2, 3, 4, 5, 6]), + codes: new Int32Array([0, 1, 2, 0, 1, 2]), + length: PROTEINS, + }; + + // Int32Array holds one code per protein, so it cannot express p1's and p4's + // second hit; the three multi-capable kinds all encode the same totals. + const MULTI_TOTALS = { A: 2, B: 2, Z: 0, [NA_VALUE]: 3 }; + const kinds: Array<[string, AnnotationData, Record]> = [ + ['Int32Array', int32, { A: 2, B: 1, Z: 0, [NA_VALUE]: 2 }], + ['sparse-multi', sparse, MULTI_TOTALS], + ['dense number[][]', dense, MULTI_TOTALS], + ['csr', csr, MULTI_TOTALS], + ]; + const alignedKinds: Array<[string, AnnotationData]> = [ + ['Int32Array', alignedInt32], + ['sparse-multi', alignedSparse], + ['dense number[][]', alignedDense], + ['csr', alignedCsr], + ]; + + it.each(kinds)( + 'matches the legacy path on %s storage (no isolation)', + (_name, colData, totals) => { + const counts = LegendDataProcessor.countFromStorage(colData, VALUES, PROTEINS, null, [ + 'A', + 'B', + 'Z', + ]); + expect(asObject(counts)).toEqual( + asObject(legacyCounts(colData, VALUES, PROTEINS, null, ['A', 'B', 'Z'])), + ); + // Explicit shape so a change to BOTH paths cannot pass unnoticed. + expect(asObject(counts)).toEqual(totals); + }, + ); + + it.each(alignedKinds)( + 'matches the legacy path on %s storage under isolation', + (_name, colData) => { + const filtered = new Set([0, 1, 4]); + const counts = LegendDataProcessor.countFromStorage(colData, VALUES, PROTEINS, filtered, [ + 'A', + 'B', + 'Z', + ]); + expect(asObject(counts)).toEqual( + asObject(legacyCounts(colData, VALUES, PROTEINS, filtered, ['A', 'B', 'Z'])), + ); + expect(asObject(counts)).toEqual({ A: 1, B: 2, Z: 0 }); + }, + ); + + it('seeds knownValues at zero and drops values nothing points at', () => { + const counts = LegendDataProcessor.countFromStorage( + new Int32Array([0, 0]), + ['A', 'B'], + 2, + null, + ['A', 'B'], + ); + expect(asObject(counts)).toEqual({ A: 2, B: 0 }); + + const unseeded = LegendDataProcessor.countFromStorage( + new Int32Array([0, 0]), + ['A', 'B'], + 2, + null, + ); + expect(asObject(unseeded)).toEqual({ A: 2 }); + }); + + it('merges distinct value slots that share an internal key', () => { + // Two bins both labelled null collapse onto the single '__NA__' key. + const counts = LegendDataProcessor.countFromStorage( + new Int32Array([0, 1, 2]), + [null, null, 'A'], + 3, + null, + ); + expect(asObject(counts)).toEqual({ [NA_VALUE]: 2, A: 1 }); + }); + + it('ignores proteins past the end of the storage', () => { + const counts = LegendDataProcessor.countFromStorage(new Int32Array([0]), ['A'], 5, null); + expect(asObject(counts)).toEqual({ A: 1 }); + }); + + // ─── isolation misalignment the flat-array path had ───────────────────── + // The legacy list held one entry per HIT, not per protein, while + // filteredIndices holds PROTEIN indices. Any storage that compacts (a + // protein with no annotation) or expands (a multi-valued protein) shifted + // the two apart, so isolation counted the wrong proteins. + + it('counts the isolated protein, not the shifted one, when codes compact', () => { + // The numeric-binning shape: -1 marks an unbinned protein + // (numeric-binning.ts materialises it that way) and the legacy build + // dropped it, shifting every later entry down by one. + const bins: (string | null)[] = ['low', 'high']; + const colData = new Int32Array([-1, 0, 1, 0]); + const isolated = new Set([0, 1]); // proteins p0 (unbinned) and p1 ('low') + + expect(asObject(LegendDataProcessor.countFromStorage(colData, bins, 4, isolated))).toEqual({ + low: 1, + }); + // Legacy list was ['low', 'high', 'low'] for proteins p1, p2, p3, so index + // 1 resolved to p2's 'high' -- a protein that is not isolated at all. + expect(asObject(legacyCounts(colData, bins, 4, isolated))).toEqual({ low: 1, high: 1 }); + }); + + it('counts every label of an isolated multi-valued protein', () => { + const labels: (string | null)[] = ['A', 'B']; + const colData: readonly (readonly number[])[] = [[0, 1], [0], [1]]; + const isolated = new Set([0]); // protein p0, which carries both labels + + expect(asObject(LegendDataProcessor.countFromStorage(colData, labels, 3, isolated))).toEqual({ + A: 1, + B: 1, + }); + // Legacy list was ['A', 'B', 'A', 'B'], so index 0 kept only 'A' and p0's + // second label was attributed to a protein that was filtered out. + expect(asObject(legacyCounts(colData, labels, 3, isolated))).toEqual({ A: 1 }); + }); + }); + describe('countAnnotationFrequencies', () => { it('counts all values when not in isolation mode', () => { const values = ['a', 'b', 'a', 'c', 'a']; diff --git a/packages/core/src/components/legend/legend-data-processor.ts b/packages/core/src/components/legend/legend-data-processor.ts index d33d382c..b05aa961 100644 --- a/packages/core/src/components/legend/legend-data-processor.ts +++ b/packages/core/src/components/legend/legend-data-processor.ts @@ -1,8 +1,11 @@ +import type { AnnotationData } from '@protspace/utils'; +import { isCsrAnnotationData, isSparseMultiValueAnnotationData } from '@protspace/utils'; import type { LegendItem, OtherItem, LegendSortMode, PersistedCategoryData } from './types'; import { getVisualEncoding, SlotTracker } from './visual-encoding'; import { LEGEND_VALUES, NA_DEFAULT_COLOR, + NA_VALUE, toInternalValue, isNAValue, toDisplayValue, @@ -72,17 +75,108 @@ export class LegendDataProcessor { return filtered; } + /** + * Count legend frequencies straight out of the annotation storage, keyed by + * PROTEIN index. + * + * Replaces "materialise a flat `string[]` of one interned label per hit, then + * reduce it to a Map". That array was a pure intermediate (at 573K proteins, + * 33 ms single-valued and 140 ms with two hits each, against 1 to 3 ms for + * this bincount) and it also broke isolation filtering: it + * compacted proteins with no annotation away and expanded multi-valued ones, + * so its indices no longer lined up with the protein indices that + * `filteredIndices` holds. Counting per protein removes both problems. + * + * `filteredIndices === null` means "no isolation filter". Per-code semantics + * match the accessor path exactly: a negative code on single-valued storage + * means "no annotation" and contributes nothing, while any code outside + * `[0, values.length)` inside a hit list resolves to `undefined` and lands in + * the `__NA__` bucket, exactly as `toInternalValue(values[code])` did. + */ + static countFromStorage( + colData: AnnotationData, + values: (string | null)[], + proteinCount: number, + filteredIndices: ReadonlySet | null, + knownValues: string[] = [], + ): Map { + const valueCount = values.length; + // One extra slot for every code that does not address a real value. + const naBin = valueCount; + const bins = new Int32Array(valueCount + 1); + + if (isSparseMultiValueAnnotationData(colData)) { + const { base, overrides } = colData; + for (let i = 0; i < proteinCount; i++) { + if (filteredIndices !== null && !filteredIndices.has(i)) continue; + const override = overrides.get(i); + if (override) { + for (let j = 0; j < override.length; j++) { + const code = override[j]; + bins[code >= 0 && code < valueCount ? code : naBin]++; + } + continue; + } + if (i >= base.length) continue; + const code = base[i]; + if (code >= 0) bins[code < valueCount ? code : naBin]++; + } + } else if (isCsrAnnotationData(colData)) { + const { end, codes } = colData; + const limit = Math.min(proteinCount, colData.length); + for (let i = 0; i < limit; i++) { + if (filteredIndices !== null && !filteredIndices.has(i)) continue; + const stop = end[i]; + for (let j = i === 0 ? 0 : end[i - 1]; j < stop; j++) { + const code = codes[j]; + bins[code >= 0 && code < valueCount ? code : naBin]++; + } + } + } else if (colData instanceof Int32Array) { + const limit = Math.min(proteinCount, colData.length); + for (let i = 0; i < limit; i++) { + if (filteredIndices !== null && !filteredIndices.has(i)) continue; + const code = colData[i]; + if (code >= 0) bins[code < valueCount ? code : naBin]++; + } + } else { + const limit = Math.min(proteinCount, colData.length); + for (let i = 0; i < limit; i++) { + if (filteredIndices !== null && !filteredIndices.has(i)) continue; + const row = colData[i]; + for (let j = 0; j < row.length; j++) { + const code = row[j]; + bins[code >= 0 && code < valueCount ? code : naBin]++; + } + } + } + + const freq = new Map(knownValues.map((value) => [value, 0] as const)); + for (let i = 0; i < valueCount; i++) { + if (bins[i] === 0) continue; + const key = toInternalValue(values[i]); + freq.set(key, (freq.get(key) ?? 0) + bins[i]); + } + if (bins[naBin] > 0) freq.set(NA_VALUE, (freq.get(NA_VALUE) ?? 0) + bins[naBin]); + return freq; + } + /** * Count frequencies of annotation values. * Raw null/empty values are converted to NA_VALUE. + * + * A frequency map (from {@link countFromStorage}) short-circuits: it is + * already the answer. */ static countAnnotationFrequencies( - annotationValues: (string | null)[], + annotationValues: (string | null)[] | ReadonlyMap, isolationMode: boolean, isolationHistory: string[][], filteredIndices: Set, knownValues: string[] = [], - ): Map { + ): ReadonlyMap { + if (!Array.isArray(annotationValues)) return annotationValues; + const freq = new Map(knownValues.map((value) => [value, 0] as const)); const countValue = (rawValue: string | null) => { @@ -108,7 +202,7 @@ export class LegendDataProcessor { * All values are internal format (N/A is '__NA__'). */ static sortAndLimitItems( - frequencyMap: Map, + frequencyMap: ReadonlyMap, maxVisibleValues: number, isolationMode: boolean, sortMode: LegendSortMode, @@ -396,7 +490,7 @@ export class LegendDataProcessor { static processLegendItems( ctx: LegendProcessorContext, annotationName: string, - annotationValues: (string | null)[], + annotationValues: (string | null)[] | ReadonlyMap, proteinIds: string[], maxVisibleValues: number, isolationMode: boolean, @@ -414,14 +508,17 @@ export class LegendDataProcessor { ): { legendItems: LegendItem[]; otherItems: OtherItem[] } { this.resetIfAnnotationChanged(ctx, annotationName); - const filteredIndices = this.getFilteredIndices(isolationMode, isolationHistory, proteinIds); - const frequencyMap = this.countAnnotationFrequencies( - annotationValues, - isolationMode, - isolationHistory, - filteredIndices, - knownValues, - ); + // A pre-counted map already has isolation applied, so the O(N) protein-id + // scan behind getFilteredIndices is skipped with it. + const frequencyMap = Array.isArray(annotationValues) + ? this.countAnnotationFrequencies( + annotationValues, + isolationMode, + isolationHistory, + this.getFilteredIndices(isolationMode, isolationHistory, proteinIds), + knownValues, + ) + : annotationValues; // Build existing zOrder map for manual sorting const existingZOrders = new Map(); diff --git a/packages/core/src/components/legend/legend.ts b/packages/core/src/components/legend/legend.ts index e09b5c64..22bcc9c7 100644 --- a/packages/core/src/components/legend/legend.ts +++ b/packages/core/src/components/legend/legend.ts @@ -33,7 +33,7 @@ import { type EatReliabilityState, type CategoryScore, } from '@protspace/utils'; -import type { LegendSettingsMap } from '@protspace/utils'; +import type { AnnotationData, LegendSettingsMap } from '@protspace/utils'; // Configuration and styles import { @@ -76,7 +76,6 @@ import { isolateItem, computeOtherConcreteValues, } from './legend-helpers'; -import { buildAnnotationValueList } from './annotation-values'; import { computeEatPopulationCounts, type EatPopulationCounts } from './eat-population-counts'; // Dialogs @@ -288,6 +287,19 @@ export class ProtspaceLegend extends LitElement { prev === undefined || !isSameReliability(next, prev), }) private _reliability: EatReliabilityState = DEFAULT_EAT_RELIABILITY; + /** + * The annotation storage the legend counts from, captured on every scatterplot + * data change. Counting straight out of it (`countFromStorage`) replaces the + * flat `annotationValues` array, which cost one interned string per protein + * and misaligned isolation filtering. `null` means "no synced storage": the + * `autoSync === false` embedding path, which still feeds the public + * `annotationValues` property instead. + */ + private _countSource: { + colData: AnnotationData; + values: (string | null)[]; + proteinCount: number; + } | null = null; @state() private _keyboardDragValue: string | null = null; private _announceManualPromotionOnNextReorder = false; private _keyboardReorderSnapshot: { @@ -1025,6 +1037,16 @@ export class ProtspaceLegend extends LitElement { this._syncNumericSettingsFromPersistence(); } + // An externally fed annotationValues array (the autoSync === false embedding + // path) is the source of truth while it is being fed, so it drops any storage + // captured by an earlier sync. Only a non-empty assignment counts: Lit reports + // the declared `= []` initializer as a change on the very first update, which + // would otherwise discard the storage the sync just captured. Clearing is + // `clearAllState`'s job. + if (changedProperties.has('annotationValues') && this.annotationValues.length > 0) { + this._countSource = null; + } + // Update legend items when relevant properties change if ( changedProperties.has('data') || @@ -1173,6 +1195,7 @@ export class ProtspaceLegend extends LitElement { this.selectedAnnotation = ''; this.annotationData = { name: '', values: [] }; this.annotationValues = []; + this._countSource = null; this.proteinIds = []; this.requestUpdate(); @@ -1627,9 +1650,34 @@ export class ProtspaceLegend extends LitElement { } private _updateAnnotationValues(data: ScatterplotData, selectedAnnotation: string): void { - const colData = data.annotation_data[selectedAnnotation]; - const values = data.annotations[selectedAnnotation].values; - this.annotationValues = buildAnnotationValueList(colData, values, data.protein_ids.length); + this._countSource = { + colData: data.annotation_data[selectedAnnotation], + values: data.annotations[selectedAnnotation].values, + proteinCount: data.protein_ids.length, + }; + } + + /** + * Legend counts for the synced storage, or `null` when there is none and the + * public `annotationValues` array is the source. + * + * Recomputed per rebuild rather than cached: the bincount is ~2 ms at 573K, + * cheaper than the array scan it replaces, and it depends on the isolation + * state, which changes independently of the storage. + */ + private _computeAnnotationCounts(knownValues: string[]): ReadonlyMap | null { + const source = this._countSource; + if (!source) return null; + const isolating = this.isolationMode && this.isolationHistory?.length > 0; + return LegendDataProcessor.countFromStorage( + source.colData, + source.values, + source.proteinCount, + isolating + ? LegendDataProcessor.getFilteredIndices(true, this.isolationHistory, this.proteinIds) + : null, + knownValues, + ); } private _hasSelectedEatAnnotation(): boolean { @@ -1770,9 +1818,14 @@ export class ProtspaceLegend extends LitElement { // Aligned with PersistenceController's isNumericAnnotation callback so the // processor and the persistence layer agree on numeric-ness in transient states. const isNumericAnnotation = this._isCurrentAnnotationNumeric(); + const knownValues = + isNumericAnnotation && this.annotationData?.values?.length + ? this.annotationData.values.map((value) => toInternalValue(value)) + : []; + const frequencies = this._computeAnnotationCounts(knownValues); if ( !this.annotationData?.values?.length || - (!isNumericAnnotation && !this.annotationValues?.length) + (!isNumericAnnotation && !(frequencies?.size ?? this.annotationValues?.length)) ) { this._legendItems = []; return; @@ -1801,9 +1854,6 @@ export class ProtspaceLegend extends LitElement { : new Set(); const numericOrderValues = this._getNumericOrderValues(); const numericDisplayLabels = this._getNumericDisplayLabelMap(); - const knownValues = isNumericAnnotation - ? this.annotationData.values.map((value) => toInternalValue(value)) - : []; const numericManualOrderIds = isNumericAnnotation ? (this._buildNumericManualOrderIds(this.selectedAnnotation) ?? []) : []; @@ -1825,7 +1875,7 @@ export class ProtspaceLegend extends LitElement { const { legendItems, otherItems } = LegendDataProcessor.processLegendItems( this._processorContext, this.annotationData.name || this.selectedAnnotation, - this.annotationValues, + frequencies ?? this.annotationValues, this.proteinIds, this.maxVisibleValues, this.isolationMode, diff --git a/packages/core/src/components/scatter-plot/scatter-plot.filter-render.test.ts b/packages/core/src/components/scatter-plot/scatter-plot.filter-render.test.ts index 9716d323..f694ff29 100644 --- a/packages/core/src/components/scatter-plot/scatter-plot.filter-render.test.ts +++ b/packages/core/src/components/scatter-plot/scatter-plot.filter-render.test.ts @@ -18,7 +18,7 @@ import { vi, describe, it, expect, afterEach } from 'vitest'; import type { PlotData, PlotDataPoint, VisualizationData } from '@protspace/utils'; import { plotDataId, materializePlotDataPoint, clonePlotData } from '@protspace/utils'; -import { buildAnnotationValueList } from '../legend/annotation-values'; +import { LegendDataProcessor } from '../legend/legend-data-processor'; vi.hoisted(() => { if (!('ResizeObserver' in globalThis)) { @@ -501,14 +501,15 @@ describe('scatter-plot data-change dispatch reflects the filtered view', () => { // Sliced to the 3 filtered ids — the legend reflects what is shown. expect(payload.protein_ids).toEqual(['p0', 'p1', 'p2']); - // The dispatched value list reflects the filtered view: only 'A' remains. - const values = buildAnnotationValueList( + // The dispatched counts reflect the filtered view: only 'A' remains. + const counts = LegendDataProcessor.countFromStorage( payload.annotation_data.fam, payload.annotations.fam.values, payload.protein_ids.length, + null, ); - expect(values).toContain('A'); - expect(values).not.toContain('B'); + expect(counts.get('A')).toBe(3); + expect(counts.has('B')).toBe(false); }); it('still slices the dispatched payload to the isolated survivors (isolation unaffected)', () => { From d665f33a49fc711523d758e79eef2afad9b5ce8d Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 01:30:40 +0200 Subject: [PATCH 14/31] feat(bundle): add the parquetbundle v3 decoder CSR counts, dictionaries and wide projections back to v2-shaped tables. --- .../src/protspace/data/io/bundle_v3.py | 255 +++++++++++ apps/protspace/tests/test_bundle_v3_decode.py | 422 ++++++++++++++++++ 2 files changed, 677 insertions(+) create mode 100644 apps/protspace/tests/test_bundle_v3_decode.py diff --git a/apps/protspace/src/protspace/data/io/bundle_v3.py b/apps/protspace/src/protspace/data/io/bundle_v3.py index b8010bb2..136ad49e 100644 --- a/apps/protspace/src/protspace/data/io/bundle_v3.py +++ b/apps/protspace/src/protspace/data/io/bundle_v3.py @@ -40,8 +40,10 @@ from protspace.data.annotations.encoding import ( FORMAT_VERSION_KEY, decode_field, + encode_field, migrate_legacy_annotation_table, read_format_version, + stamp_format_version, ) CONTAINER_VERSION = 3 @@ -549,3 +551,256 @@ def encode_v3( _write(projections_table), _write(payload_table), ) + + +# --------------------------------------------------------------------------- # +# decoder +# --------------------------------------------------------------------------- # + + +def _read(part: bytes) -> pa.Table: + return pq.read_table(io.BytesIO(part)) + + +def _flat(column: pa.ChunkedArray | pa.Array) -> pa.Array: + """One contiguous Arrow array (``ListArray.from_arrays`` refuses chunks).""" + if not isinstance(column, pa.ChunkedArray): + return column + if column.num_chunks == 1: + return column.chunk(0) + if column.num_chunks == 0: + return pa.array([], type=column.type) + return pa.concat_arrays(column.chunks) + + +def _read_payloads(part: bytes) -> dict[str, bytes]: + table = _read(part) + return dict( + zip( + table.column("name").to_pylist(), + table.column("data").to_pylist(), + strict=True, + ) + ) + + +def _read_labels(payloads: dict[str, bytes], name: str) -> list[str]: + """Slice ``dict:`` by the prefix sum of its per-label byte lengths.""" + blob = payloads[f"dict:{name}"] + lengths = np.frombuffer(payloads[f"dict:{name}:len"], " pa.Array: + """Prefix-sum per-element ``counts`` into list offsets, then join each list.""" + offsets = np.concatenate(([0], np.cumsum(counts, dtype=np.int64))).astype(np.int32) + lists = pa.ListArray.from_arrays(pa.array(offsets, type=pa.int32()), values) + return pc.binary_join(lists, separator) + + +def _restorable_type(alias: str) -> pa.DataType | None: + """The numeric Arrow type ``alias`` names, or *None* to render v2 strings. + + A string ``sourceType`` deliberately lands here: the v2 spelling of the + column is what the encoder consumed, so rendering it back is the restoration. + """ + if alias in (_UNRESTORABLE_SOURCE_TYPE, "string", "large_string"): + return None + try: + type_ = pa.type_for_alias(alias) + except ValueError: + return None + return type_ if pa.types.is_integer(type_) or pa.types.is_floating(type_) else None + + +def _decode_numeric(column: pa.ChunkedArray, entry: dict[str, Any]) -> pa.Array: + """float64 + NaN back to the source Arrow type, or to its v2 string cells.""" + values = _flat(column).to_numpy(zero_copy_only=False) + present = ~np.isnan(values) + type_ = _restorable_type(entry.get("sourceType", _UNRESTORABLE_SOURCE_TYPE)) + if type_ is not None: + return pc.cast(pa.array(values, mask=~present), type_) + + finite = np.where(present, values, 0.0) + # ``str(2.0)`` is ``"2.0"`` but an int-typed v2 column spells it ``"2"``, and + # numpy's float repr is Python's, so int columns take the int64 detour. The + # magnitude guard keeps a value past int64 out of an undefined cast. + if entry.get("numericType") == "int" and np.abs(finite).max(initial=0.0) < 2.0**63: + text = finite.astype(np.int64).astype(str) + else: + text = finite.astype(str) + return pc.if_else( + pa.array(present), pa.array(text, type=pa.string()), pa.scalar("") + ) + + +def _decode_categorical(column: pa.ChunkedArray, labels: pa.Array) -> pa.Array: + """int32 codes back to label cells; ``-1`` (missing) becomes ``""``.""" + codes = _flat(column).to_numpy(zero_copy_only=False) + return pc.fill_null(labels.take(pa.array(codes, mask=codes < 0)), "") + + +def _decode_multi( + column: pa.ChunkedArray, + name: str, + entry: dict[str, Any], + payloads: dict[str, bytes], + labels: pa.Array, + evidence_labels: pa.Array, +) -> pa.Array: + """CSR hits back to ``label|suffix;label|suffix`` cells (``""`` when empty).""" + hits = labels.take(pa.array(np.frombuffer(payloads[f"csr:{name}"], " 0), + _list_join(per_hit, text, ","), + pa.scalar(None, pa.string()), + ) + suffix = scored if suffix is None else pc.coalesce(suffix, scored) + + if suffix is not None: + # A null suffix (no evidence, no scores) leaves the bare label: the + # element-wise join emits null as soon as one side is null. + hits = pc.coalesce(pc.binary_join_element_wise(hits, suffix, "|"), hits) + + counts = _flat(column).to_numpy(zero_copy_only=False) + return _list_join(counts, hits, ";") + + +def _decode_projections( + part: bytes, manifest: list[dict[str, Any]], identifiers: pa.Array +) -> pa.Table: + """Wide float32 projections back to the long v2 table, in manifest order.""" + wide = _read(part) + num_rows = len(identifiers) + row = pa.array(np.zeros(num_rows, dtype=np.int32)) + schema = pa.schema( + [ + ("projection_name", pa.string()), + ("identifier", pa.string()), + ("x", pa.float32()), + ("y", pa.float32()), + ("z", pa.float32()), + ] + ) + + tables = [] + for projection in manifest: + name = projection["name"] + dimension = int(projection["dimension"]) + tables.append( + pa.table( + { + # ``take`` of a one-element array beats materialising N copies. + "projection_name": pa.array([name], type=pa.string()).take(row), + "identifier": identifiers, + "x": _flat(wide.column(f"{name}__x")), + "y": _flat(wide.column(f"{name}__y")), + "z": _flat(wide.column(f"{name}__z")) + if dimension == 3 + else pa.nulls(num_rows, pa.float32()), + }, + schema=schema, + ) + ) + return pa.concat_tables(tables) if tables else schema.empty_table() + + +def decode_v3(parts: list[bytes]) -> tuple[pa.Table, pa.Table, pa.Table]: + """Decode v3 parts back into the three v2-shaped tables. + + ``parts`` is what :func:`encode_v3` returned: annotations, projections + metadata, wide projections, payloads (bundle parts 1, 2, 3 and 6). The + result is re-stamped ``protspace_format_version=2`` because what comes back + *is* the v2 cell grammar every Python consumer parses. + + The round trip is not byte-exact, and deliberately so -- v3 stores what the + browser's v2 reader would have parsed out of the cells, not the cells: + + * hits and cells are whitespace-trimmed, and empty or missing-valued hits are + dropped (``"A;;B"`` comes back ``"A;B"``, ``" A |IDA"`` as ``"A|IDA"``); + * a missing cell -- null, blank or a ``MISSING_TOKENS`` spelling -- comes back + as ``""``; + * labels are re-encoded canonically, so ``%3b`` comes back as ``%3B``; + * scores round-trip through float32 and are re-spelled shortest-first, so + ``"0.5700"`` comes back as ``"0.57"``; + * a numeric column comes back in its ``sourceType`` when that is restorable + and otherwise as its canonical v2 spelling, so an all-integral column + spells ``100``, never ``100.0``; + * projection coordinates come back float32 (``z`` null for a 2D projection) + and a protein absent from a projection comes back at the origin; + * the identifier column comes back first, wherever it sat before. + """ + if len(parts) != 4: + raise ValueError( + f"decode_v3 expects the 4 parts encode_v3 returns, got {len(parts)}" + ) + + annotations = _read(parts[0]) + metadata = dict(annotations.schema.metadata or {}) + raw_manifest = metadata.pop(MANIFEST_KEY, None) + if raw_manifest is None: + raise ValueError( + f"annotations part carries no {MANIFEST_KEY.decode()} key; " + "it is not a v3 part" + ) + manifest = json.loads(raw_manifest) + payloads = _read_payloads(parts[3]) + + evidence_labels = pa.array( + _read_labels(payloads, _EVIDENCE_DICT_NAME) + if f"dict:{_EVIDENCE_DICT_NAME}" in payloads + else [], + type=pa.string(), + ) + + id_column = manifest["idColumn"] + columns: dict[str, pa.Array] = {id_column: _flat(annotations.column(id_column))} + for name, entry in manifest["columns"].items(): + kind = entry["kind"] + if kind == "numeric": + columns[name] = _decode_numeric(annotations.column(name), entry) + continue + # Labels are stored decoded; the v2 cell grammar wants them encoded. + labels = pa.array( + [encode_field(label) for label in _read_labels(payloads, name)], + type=pa.string(), + ) + if kind == "categorical": + columns[name] = _decode_categorical(annotations.column(name), labels) + elif kind == "multi": + columns[name] = _decode_multi( + annotations.column(f"{name}__count"), + name, + entry, + payloads, + labels, + evidence_labels, + ) + else: + raise ValueError(f"column '{name}' has unknown v3 kind '{kind}'") + + return ( + stamp_format_version(pa.table(columns).replace_schema_metadata(metadata)), + _read(parts[1]), + _decode_projections(parts[2], manifest["projections"], columns[id_column]), + ) diff --git a/apps/protspace/tests/test_bundle_v3_decode.py b/apps/protspace/tests/test_bundle_v3_decode.py new file mode 100644 index 00000000..1703b62e --- /dev/null +++ b/apps/protspace/tests/test_bundle_v3_decode.py @@ -0,0 +1,422 @@ +"""Decoder half of parquetbundle format v3 (``data/io/bundle_v3.decode_v3``). + +Six Python consumers (``utils/arrow_reader``, ``cli/serve``, ``cli/style`` + +``utils/add_annotation_style``, ``cli/transfer``, ``cli/bundle`` and the +``scripts/``) parse the v2 string grammar, so v3 only ever exists between +``write_bundle`` and ``read_tables``. The contract these tests pin is therefore +``decode_v3(encode_v3(T)) == T`` on pipeline-shaped tables, plus the handful of +places where that equality is deliberately *not* exact: v3 stores what the +browser's v2 reader would have parsed out of a cell, not the cell. +""" + +import io +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from protlabel import Prediction +from protspace.data.annotations.encoding import ( + FORMAT_VERSION_KEY, + stamp_format_version, +) +from protspace.data.io.bundle_v3 import MANIFEST_KEY, decode_v3, encode_v3 +from protspace.data.io.predictions import add_overlay_columns +from protspace.data.processors.base_processor import BaseProcessor + +REAL_BUNDLE = ( + Path(__file__).resolve().parents[3] + / "apps" + / "web" + / "public" + / "data" + / "venom_eat_stats.parquetbundle" +) + + +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # + + +def annotations_table(**columns: list[str]) -> pa.Table: + """An annotations table exactly as the pipeline builds it (all-string, v2).""" + n = len(next(iter(columns.values()))) + frame = pd.DataFrame({"identifier": [f"p{i}" for i in range(n)], **columns}) + return BaseProcessor({}, {})._create_protein_annotations_table(frame) + + +def projection_tables(num_rows: int, dimensions=(2, 3)): + """``projections_metadata`` + long ``projections_data`` for ``num_rows`` proteins.""" + processor = BaseProcessor({}, {}) + reductions = [ + { + "name": f"PCA {dimension}", + "dimensions": dimension, + "info": {"components": dimension}, + "data": np.arange(num_rows * dimension, dtype=np.float32).reshape( + num_rows, dimension + ), + } + for dimension in dimensions + ] + headers = [f"p{i}" for i in range(num_rows)] + return ( + processor._create_projections_metadata_table(reductions), + processor._create_projections_data_table(reductions, headers), + ) + + +def round_trip(annotations: pa.Table, dimensions=(2, 3)): + """Encode then decode ``annotations`` with matching projections.""" + metadata, data = projection_tables(annotations.num_rows, dimensions) + return decode_v3(encode_v3(annotations, metadata, data)) + + +def cells(annotations: pa.Table, column: str, dimensions=(2, 3)) -> list: + return round_trip(annotations, dimensions)[0].column(column).to_pylist() + + +# --------------------------------------------------------------------------- # +# the round trip on pipeline-shaped tables +# --------------------------------------------------------------------------- # + + +def pipeline_annotations() -> pa.Table: + """Every cell shape the encoder dispatches on, in one pipeline-built table.""" + table = annotations_table( + # plain categorical, and one all-empty column + kingdom=["Bacteria", "Archaea", "Bacteria", "Eukaryota", "Archaea", "Bacteria"], + unknown=["", "", "", "", "", ""], + # multi + scores, with zero hits first, interior and last, and labels + # carrying the encoded ';' and '|' the v2 grammar reserves + pfam=[ + "", + "PF00001 (7tm%3B1)|1e-10,2.5;PF00002|0.5", + "", + "PF00001 (7tm%3B1)|0.25", + "PF00003 (a%7Cb)|3;PF00004", + "", + ], + # multi + evidence + go_mf=[ + "GO:0005524|IDA", + "", + "GO:0005524|IDA;GO:0016787|ECO:0000269", + "GO:0016787|IEA", + "", + "GO:0005524|IDA", + ], + # numeric int with blanks, numeric float + length=["100", "", "250", "", "3000", "42"], + annotation_score=["0.5", "1.25", "", "0.5", "2.0", "0.125"], + ) + return overlay(table, range(table.num_rows)) + + +def overlay(table: pa.Table, predicted) -> pa.Table: + """Attach EAT ``ec__pred_*`` companions for the given row indices.""" + predictions = [ + Prediction( + query_id=f"p{i}", + label="3.4.21.- (Serine endopeptidases)", + source_id="P20005", + distance=0.5, + reliability=0.35313386, + k=1, + metric="euclidean", + ) + for i in predicted + ] + return add_overlay_columns( + table, + "ec", + predictions, + identifiers=table.column("protein_id").to_pylist(), + ) + + +def test_pipeline_round_trip_is_cell_for_cell(): + source = pipeline_annotations() + metadata, data = projection_tables(source.num_rows) + decoded, decoded_metadata, decoded_data = decode_v3( + encode_v3(source, metadata, data) + ) + + assert decoded.column_names == source.column_names + for name in source.column_names: + assert decoded.column(name).to_pylist() == source.column(name).to_pylist(), name + assert decoded.equals(source) + + assert decoded_metadata.equals(metadata) + assert decoded_data.to_pydict() == data.to_pydict() + + +def test_a_partial_eat_overlay_keeps_null_confidences_but_blanks_the_strings(): + """float32 nulls survive ``sourceType``; string nulls hit the one missing code.""" + source = overlay(annotations_table(kingdom=["A", "B"]), [1]) + decoded = round_trip(source, (2,))[0] + assert decoded.schema.field("ec__pred_confidence").type == pa.float32() + assert decoded.column("ec__pred_confidence").to_pylist() == [ + None, + pytest.approx(0.35313386), + ] + assert decoded.column("ec__pred_source").to_pylist() == ["", "P20005"] + + +def test_footer_says_two_and_the_manifest_key_is_gone(): + source = pipeline_annotations() + decoded = round_trip(source)[0] + assert decoded.schema.metadata[FORMAT_VERSION_KEY] == b"2" + assert MANIFEST_KEY not in decoded.schema.metadata + # ``stamp_format_version`` merges, so the pandas key the pipeline wrote lives on. + assert decoded.schema.metadata == source.schema.metadata + + +def test_decoded_fields_are_nullable_again(): + """Part 1 is written REQUIRED for hyparquet; a v2-shaped table is not.""" + decoded = round_trip(pipeline_annotations())[0] + assert all(field.nullable for field in decoded.schema) + + +# --------------------------------------------------------------------------- # +# projections +# --------------------------------------------------------------------------- # + + +def test_projections_are_long_manifest_ordered_and_protein_ordered(): + source = annotations_table(kingdom=["A", "B", "C"]) + metadata, data = projection_tables(3, (2, 3)) + _, decoded_metadata, decoded_data = decode_v3(encode_v3(source, metadata, data)) + + assert decoded_metadata.column("projection_name").to_pylist() == ["PCA 2", "PCA 3"] + columns = decoded_data.to_pydict() + assert columns["projection_name"] == ["PCA 2"] * 3 + ["PCA 3"] * 3 + assert columns["identifier"] == ["p0", "p1", "p2"] * 2 + assert columns["x"] == [0.0, 2.0, 4.0, 0.0, 3.0, 6.0] + assert columns["z"] == [None, None, None, 2.0, 5.0, 8.0] + assert decoded_data.schema.field("z").type == pa.float32() + + +def test_a_protein_absent_from_a_projection_comes_back_at_the_origin(): + """v2's zero-initialised Float32Array is the contract, not NaN.""" + source = annotations_table(kingdom=["A", "B", "C"]) + metadata, data = projection_tables(3, (2,)) + data = data.filter(pa.compute.not_equal(data.column("identifier"), pa.scalar("p1"))) + decoded_data = decode_v3(encode_v3(source, metadata, data))[2] + assert decoded_data.to_pydict()["x"] == [0.0, 0.0, 4.0] + + +# --------------------------------------------------------------------------- # +# numeric restoration +# --------------------------------------------------------------------------- # + + +def test_source_type_restores_non_string_numeric_columns(): + source = stamp_format_version( + pa.table( + { + "protein_id": ["p0", "p1", "p2"], + "length": pa.array([10, None, 30], type=pa.int32()), + "confidence": pa.array([0.5, 1.5, None], type=pa.float32()), + } + ) + ) + decoded = round_trip(source, (2,))[0] + assert decoded.schema.field("length").type == pa.int32() + assert decoded.schema.field("confidence").type == pa.float32() + assert decoded.column("length").to_pylist() == [10, None, 30] + assert decoded.column("confidence").to_pylist() == [0.5, 1.5, None] + + +def test_int_columns_never_come_back_with_a_decimal_point(): + assert cells(annotations_table(length=["100", "", "3"]), "length") == [ + "100", + "", + "3", + ] + + +def test_float_columns_keep_the_python_float_spelling(): + values = ["1.5", "2.0", "1e-10", ""] + assert cells(annotations_table(score=values), "score") == values + + +def test_an_unrestorable_source_type_falls_back_to_strings(): + """``str(dictionary<...>)`` is no alias, so the v2 spelling is the fallback.""" + source = stamp_format_version( + pa.table( + { + "protein_id": ["p0", "p1"], + "species": pa.array(["Human", "Mouse"]).dictionary_encode(), + } + ) + ) + decoded = round_trip(source, (2,))[0] + assert decoded.schema.field("species").type == pa.string() + assert decoded.column("species").to_pylist() == ["Human", "Mouse"] + + +# --------------------------------------------------------------------------- # +# deliberate non-identities +# --------------------------------------------------------------------------- # + + +def test_hits_and_cells_are_trimmed_and_empty_hits_collapse(): + table = annotations_table(pfam=[" A ;B ", "A;;B", "A; ;B"]) + assert cells(table, "pfam") == ["A;B", "A;B", "A;B"] + + +def test_missing_spellings_all_come_back_as_the_empty_string(): + table = annotations_table(col=["A", "NA", "n/a", "None", "__NA__", " "]) + assert cells(table, "col") == ["A", "", "", "", "", ""] + + +def test_null_cells_come_back_as_the_empty_string(): + """v3 has one missing code, so a null and a blank are the same cell.""" + source = stamp_format_version( + pa.table({"protein_id": ["p0", "p1"], "col": ["A", None]}) + ) + assert round_trip(source, (2,))[0].column("col").to_pylist() == ["A", ""] + + +def test_a_raw_pipe_in_a_label_comes_back_percent_encoded(): + """v2 requires ``|`` inside a label to be escaped; decode emits the legal form.""" + table = annotations_table(col=["PF3 (a|b)|0.5"]) + assert cells(table, "col") == ["PF3 (a%7Cb)|0.5"] + + +def test_percent_encoding_is_normalised_to_upper_case(): + table = annotations_table(col=["a%3bb|0.5", "a%3Bb|0.5"]) + assert cells(table, "col") == ["a%3Bb|0.5", "a%3Bb|0.5"] + + +def test_an_unscored_hit_in_a_scored_column_keeps_no_suffix(): + """``score_count`` is per hit, so a bare hit must not gain a dangling ``|``.""" + table = annotations_table(col=["PF1|0.5;PF2", "PF3"]) + assert cells(table, "col") == ["PF1|0.5;PF2", "PF3"] + + +def test_scores_round_trip_through_float32(): + table = annotations_table(col=["A|0.5700", "A|1", "A|0.1", "A|1e-10,2.5"]) + # 0.5700 loses its trailing zero (float32 has no such notion) and an integral + # score keeps the JavaScript spelling ``[1].join(',') === '1'``. + assert cells(table, "col") == ["A|0.57", "A|1", "A|0.1", "A|1e-10,2.5"] + + +def test_an_int_column_re_spells_its_cells_canonically(): + table = annotations_table(col=["1", "2.0", "+3", "4e1"]) + assert cells(table, "col") == ["1", "2", "3", "40"] + + +def test_a_bool_column_comes_back_as_the_python_spelling(): + """``sourceType`` restoration is numeric-only; a bool stays v2's ``True``/``False``.""" + source = stamp_format_version( + pa.table({"protein_id": ["p0", "p1"], "flag": [True, False]}) + ) + decoded = round_trip(source, (2,))[0] + assert decoded.schema.field("flag").type == pa.string() + assert decoded.column("flag").to_pylist() == ["True", "False"] + + +# --------------------------------------------------------------------------- # +# guards +# --------------------------------------------------------------------------- # + + +def test_rejects_a_part_list_that_is_not_the_encoder_output(): + with pytest.raises(ValueError, match="expects the 4 parts"): + decode_v3([b"", b"", b""]) + + +def test_rejects_an_annotations_part_without_a_manifest(): + source = annotations_table(col=["A", "B"]) + parts = list(encode_v3(source, *projection_tables(2, (2,)))) + without = pq.read_table(io.BytesIO(parts[0])).replace_schema_metadata( + {FORMAT_VERSION_KEY: b"3"} + ) + buffer = io.BytesIO() + pq.write_table(without, buffer) + parts[0] = buffer.getvalue() + with pytest.raises(ValueError, match="not a v3 part"): + decode_v3(parts) + + +def test_rejects_an_unknown_kind(): + source = annotations_table(col=["A", "B"]) + parts = list(encode_v3(source, *projection_tables(2, (2,)))) + table = pq.read_table(io.BytesIO(parts[0])) + manifest = json.loads(table.schema.metadata[MANIFEST_KEY]) + manifest["columns"]["col"]["kind"] = "sparse" + buffer = io.BytesIO() + pq.write_table( + table.replace_schema_metadata( + {**table.schema.metadata, MANIFEST_KEY: json.dumps(manifest).encode()} + ), + buffer, + ) + parts[0] = buffer.getvalue() + with pytest.raises(ValueError, match="unknown v3 kind"): + decode_v3(parts) + + +# --------------------------------------------------------------------------- # +# the real shipped bundle +# --------------------------------------------------------------------------- # + + +@pytest.mark.skipif(not REAL_BUNDLE.exists(), reason="web sample data not checked out") +def test_real_bundle_round_trip(): + """``venom_eat_stats`` (v2, 811 x 38) end to end, with its non-identities named. + + Every column that is not byte-identical is one of the two documented losses, + and nothing else drifts: 4 cluster columns whose ``%.4f`` scores lose a + trailing zero to float32, and 4 all-or-partly-null overlay columns whose + nulls become ``""``. + """ + parts = REAL_BUNDLE.read_bytes().split(b"---PARQUET_DELIMITER---") + source = pq.read_table(io.BytesIO(parts[0])) + metadata = pq.read_table(io.BytesIO(parts[1])) + data = pq.read_table(io.BytesIO(parts[2])) + + decoded, decoded_metadata, decoded_data = decode_v3( + encode_v3(source, metadata, data) + ) + assert decoded.column_names == source.column_names + assert decoded.schema.metadata == source.schema.metadata + + differing = { + name + for name in source.column_names + if decoded.column(name).to_pylist() != source.column(name).to_pylist() + } + assert differing == { + "cluster_elbow_ProtT5 — PCA 2", + "cluster_silhouette_ProtT5 — PCA 2", + "cluster_elbow_ProtT5 — UMAP 2", + "cluster_silhouette_ProtT5 — UMAP 2", + "ec__pred_value", + "ec__pred_source", + "protein_families__pred_value", + "protein_families__pred_source", + } + for name in differing: + for before, after in zip( + source.column(name).to_pylist(), + decoded.column(name).to_pylist(), + strict=True, + ): + if before == after: + continue + if before is None: + assert after == "" # the null / blank collapse + else: # "cluster 4|0.5700" -> "cluster 4|0.57" + label, _, score = before.rpartition("|") + assert after == f"{label}|{float(score):g}" + + assert decoded_metadata.equals(metadata) + assert decoded_data.to_pydict() == data.to_pydict() From e45869b723ea6e08889ae30aa63b3033e3b8165c Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 01:31:44 +0200 Subject: [PATCH 15/31] perf(utils): hash datasets without BigInt and memoize the fingerprint Lane-based FNV-1a 64 plus a protein_ids-keyed memo; hash values unchanged. --- packages/utils/src/storage/data-hash.test.ts | 142 +++++++++++++++++++ packages/utils/src/storage/data-hash.ts | 126 ++++++++++++---- 2 files changed, 239 insertions(+), 29 deletions(-) diff --git a/packages/utils/src/storage/data-hash.test.ts b/packages/utils/src/storage/data-hash.test.ts index 84e2eb39..b365a93a 100644 --- a/packages/utils/src/storage/data-hash.test.ts +++ b/packages/utils/src/storage/data-hash.test.ts @@ -426,3 +426,145 @@ describe('generateDatasetHash', () => { expect(hash).toMatch(/^[0-9a-f]{16}$/); }); }); + +/** + * Hash values key localStorage for legend, control-bar and tooltip persistence, so + * the lane-based FNV must reproduce the exact strings the BigInt implementation + * produced, and the protein ordering must stay on `localeCompare`. These references + * are verbatim copies of the original implementation. + */ +function referenceFnv1a64(str: string): string { + let hash = 0xcbf29ce484222325n; + const fnvPrime = 0x100000001b3n; + const mask = 0xffffffffffffffffn; + + for (let i = 0; i < str.length; i++) { + hash ^= BigInt(str.charCodeAt(i)); + hash = (hash * fnvPrime) & mask; + } + + return hash.toString(16).padStart(16, '0'); +} + +function referenceIdOnlyHash(proteinIds: readonly string[]): string { + const order = Array.from({ length: proteinIds.length }, (_, index) => index); + order.sort((left, right) => proteinIds[left].localeCompare(proteinIds[right]) || left - right); + return referenceFnv1a64(`${order.map((index) => proteinIds[index]).join('\x00')}\x02\x02`); +} + +/** Deterministic LCG so a failure is reproducible from the seed alone. */ +function createRandom(seed: number): () => number { + let state = seed >>> 0; + return () => (state = (state * 1103515245 + 12345) >>> 0) / 4294967296; +} + +describe('generateDatasetHash value stability', () => { + it('matches the BigInt FNV-1a 64 reference on 10k random and adversarial strings', () => { + const random = createRandom(0x5eed); + const corpus = [ + '', + 'a', + '\x00', + '￿', + '\ud800', // lone high surrogate + '\udfff', // lone low surrogate + '😀', // astral pair + 'こんにちは', + 'P12345\x1fP67890', + 'x'.repeat(200_000), + ]; + for (let index = 0; index < 10_000; index++) { + const length = Math.floor(random() * 40); + let value = ''; + for (let position = 0; position < length; position++) { + value += String.fromCharCode(Math.floor(random() * 0x10000)); + } + corpus.push(value); + } + + // A single-element array joins to exactly that element, so this is the raw hash. + for (const value of corpus) { + expect(generateDatasetHash([value])).toBe(referenceFnv1a64(value)); + } + }); + + it('matches the localeCompare reference ordering on collation-sensitive ids', () => { + const random = createRandom(0xc011a); + const alphabet = [...'abcXYZ0189-_.', 'é', 'ü', '漢', 'ß', ' ', '😀']; + const corpora: string[][] = [ + ['a', 'A', 'B', 'b', 'ä', 'z', 'Z'], + ['P1', 'p1', 'P10', 'P2', 'P-1', 'P_1', 'P.1'], + ['', 'a', '', 'a'], + ['résumé', 'resume', 'RESUME', 'Résumé'], + ]; + for (let index = 0; index < 200; index++) { + const size = 1 + Math.floor(random() * 30); + corpora.push( + Array.from({ length: size }, () => { + const length = 1 + Math.floor(random() * 6); + let value = ''; + for (let position = 0; position < length; position++) { + value += alphabet[Math.floor(random() * alphabet.length)]; + } + return value; + }), + ); + } + + for (const proteinIds of corpora) { + expect(generateDatasetHash({ protein_ids: proteinIds })).toBe( + referenceIdOnlyHash(proteinIds), + ); + } + }); +}); + +describe('generateDatasetHash memoization', () => { + const buildDataset = () => ({ + protein_ids: ['P3', 'P1', 'P2'], + annotations: { length: { kind: 'numeric' as const, values: [] } }, + numeric_annotation_data: { length: [30, 10, 20] as (number | null)[] }, + annotation_predicted: { + ec: [null, { value: '1.1.1.1', confidence: 0.8, source: 'P1' }, null], + }, + }); + + it('returns the freshly computed hash on a repeat call with the same references', () => { + const dataset = buildDataset(); + const memoized = generateDatasetHash(dataset); + // Structurally identical, but every reference is new, so this cannot hit the memo. + expect(memoized).toBe(generateDatasetHash(buildDataset())); + expect(generateDatasetHash(dataset)).toBe(memoized); + }); + + it('recomputes when annotations, numeric data, or predictions are replaced', () => { + const dataset = buildDataset(); + const baseline = generateDatasetHash(dataset); + + expect( + generateDatasetHash({ + ...dataset, + numeric_annotation_data: { length: [30, 10, 21] }, + }), + ).not.toBe(baseline); + + expect( + generateDatasetHash({ + ...dataset, + annotations: { length: { kind: 'categorical' as const, values: ['a', 'b'] } }, + }), + ).not.toBe(baseline); + + expect( + generateDatasetHash({ + ...dataset, + annotation_predicted: { + ec: [null, { value: '9.9.9.9', confidence: 0.8, source: 'P1' }, null], + }, + }), + ).not.toBe(baseline); + + // The original references still resolve to the original hash. + expect(generateDatasetHash(dataset)).toBe(baseline); + }); +}); diff --git a/packages/utils/src/storage/data-hash.ts b/packages/utils/src/storage/data-hash.ts index 5f768c39..f87c3373 100644 --- a/packages/utils/src/storage/data-hash.ts +++ b/packages/utils/src/storage/data-hash.ts @@ -50,32 +50,61 @@ export function djb2Hash(str: string): number { return hash >>> 0; // Convert to unsigned 32-bit integer } -function fnv1a64Hash(str: string): string { - let hash = 0xcbf29ce484222325n; - const fnvPrime = 0x100000001b3n; - const mask = 0xffffffffffffffffn; - - for (let i = 0; i < str.length; i++) { - hash ^= BigInt(str.charCodeAt(i)); - hash = (hash * fnvPrime) & mask; - } +/** + * FNV-1a 64 carried as two 32-bit lanes instead of a BigInt. + * + * The BigInt form allocated one BigInt per character, which at 573K proteins is + * millions of allocations on the main thread. The lane form is value-identical: + * with `hash = hi * 2^32 + lo` and the prime `2^40 + 0x1b3`, + * + * hash * prime = hi*2^72 + hi*0x1b3*2^32 + lo*2^40 + lo*0x1b3 (mod 2^64) + * + * `hi*2^72` vanishes mod 2^64; `lo*2^40 mod 2^64` is `((lo << 8) >>> 0) * 2^32`; + * `lo*0x1b3` contributes its low word plus a carry into the high word. Every + * intermediate stays below 2^42, so it is exact in a double. + */ +interface Fnv1a64State { + hi: number; + lo: number; +} - return hash.toString(16).padStart(16, '0'); +function createFNV1a64(): Fnv1a64State { + return { hi: 0xcbf29ce4, lo: 0x84222325 }; } -function appendFNV1a64(hash: bigint, value: string): bigint { - const fnvPrime = 0x100000001b3n; - const mask = 0xffffffffffffffffn; - let nextHash = hash; +function appendFNV1a64(state: Fnv1a64State, value: string): void { + let hi = state.hi; + let lo = state.lo; for (let i = 0; i < value.length; i++) { - nextHash ^= BigInt(value.charCodeAt(i)); - nextHash = (nextHash * fnvPrime) & mask; + lo = (lo ^ value.charCodeAt(i)) >>> 0; + const product = lo * 0x1b3; + const nextLo = product >>> 0; + hi = (hi * 0x1b3 + (product - nextLo) / 4294967296 + ((lo << 8) >>> 0)) >>> 0; + lo = nextLo; } - return nextHash; + state.hi = hi; + state.lo = lo; } +function formatFNV1a64(state: Fnv1a64State): string { + return state.hi.toString(16).padStart(8, '0') + state.lo.toString(16).padStart(8, '0'); +} + +function fnv1a64Hash(str: string): string { + const state = createFNV1a64(); + appendFNV1a64(state, str); + return formatFNV1a64(state); +} + +/** + * Do not "optimize" this into a hoisted `new Intl.Collator()`. It is value-identical + * (ECMA-402 defines bare `localeCompare` as `new Intl.Collator(undefined, undefined) + * .compare(...)`), but V8 already caches the default collator and takes a fast path + * for one-byte strings: measured on 573K SwissProt accessions the hoisted collator + * sorts in 26 ms against 10 ms for `localeCompare`. + */ function buildProteinIndexOrder(proteinIds: readonly string[]): number[] { const order = Array.from({ length: proteinIds.length }, (_, index) => index); order.sort((left, right) => proteinIds[left].localeCompare(proteinIds[right]) || left - right); @@ -90,7 +119,7 @@ function buildNumericFingerprint( return ''; } - let hash = 0xcbf29ce484222325n; + const hash = createFNV1a64(); let nonNullCount = 0; let min = Number.POSITIVE_INFINITY; let max = Number.NEGATIVE_INFINITY; @@ -98,8 +127,8 @@ function buildNumericFingerprint( for (let position = 0; position < values.length; position++) { const value = values[proteinIndexOrder?.[position] ?? position]; const serialized = value == null ? 'null' : String(value); - hash = appendFNV1a64(hash, serialized); - hash = appendFNV1a64(hash, '\x1f'); + appendFNV1a64(hash, serialized); + appendFNV1a64(hash, '\x1f'); if (value == null || !Number.isFinite(value)) { continue; @@ -115,7 +144,7 @@ function buildNumericFingerprint( nonNullCount, nonNullCount > 0 ? min : 'none', nonNullCount > 0 ? max : 'none', - hash.toString(16).padStart(16, '0'), + formatFNV1a64(hash), ].join('|'); } @@ -177,12 +206,12 @@ function buildDatasetFingerprint(data: DatasetHashInput): string { const predictionFingerprint = Object.entries(data.annotation_predicted ?? {}) .sort(([leftName], [rightName]) => leftName.localeCompare(rightName)) .map(([annotationName, cells]) => { - let hash = 0xcbf29ce484222325n; + const hash = createFNV1a64(); let count = 0; const appendCell = (cell: PredictedCell | null, proteinId: string): void => { if (!cell) return; count += 1; - hash = appendFNV1a64( + appendFNV1a64( hash, [ proteinId, @@ -193,7 +222,7 @@ function buildDatasetFingerprint(data: DatasetHashInput): string { cell.source, ].join('\x1f'), ); - hash = appendFNV1a64(hash, '\x1e'); + appendFNV1a64(hash, '\x1e'); }; for (let index = proteinIds.length; index < cells.length; index++) { appendCell(cells[index], ''); @@ -201,21 +230,60 @@ function buildDatasetFingerprint(data: DatasetHashInput): string { for (const index of proteinIndexOrder) { appendCell(cells[index] ?? null, proteinIds[index]); } - return `${annotationName}::${count}::${hash.toString(16).padStart(16, '0')}`; + return `${annotationName}::${count}::${formatFNV1a64(hash)}`; }) .join('\x01'); return [sortedIds.join('\x00'), annotationFingerprint, predictionFingerprint].join('\x02'); } +/** + * Callers rebuild the wrapper object on every data change (`legend.ts`) while the + * inner arrays keep their identity, so keying on `protein_ids` identity plus the + * three payload references turns the repeat calls into a lookup. This is only + * sound because no producer mutates a live dataset in place: every transform + * (`materializeVisualizationData`, `cloneWithPredictions`, the conversion + * pipeline) hands back fresh containers, which miss the memo and recompute. + */ +interface DatasetHashMemo { + annotations: DatasetHashInput['annotations']; + numericAnnotationData: DatasetHashInput['numeric_annotation_data']; + annotationPredicted: DatasetHashInput['annotation_predicted']; + hash: string; +} + +const datasetHashMemo = new WeakMap(); + export function generateDatasetHash(input: string[] | DatasetHashInput): string { if (!input || (Array.isArray(input) && input.length === 0)) { return '0000000000000000'; } - const combined = Array.isArray(input) - ? [...input].sort().join('\x00') - : buildDatasetFingerprint(input); + if (Array.isArray(input)) { + return fnv1a64Hash([...input].sort().join('\x00')); + } + + const memoKey = Array.isArray(input.protein_ids) ? input.protein_ids : null; + const memo = memoKey ? datasetHashMemo.get(memoKey) : undefined; + if ( + memo && + memo.annotations === input.annotations && + memo.numericAnnotationData === input.numeric_annotation_data && + memo.annotationPredicted === input.annotation_predicted + ) { + return memo.hash; + } + + const hash = fnv1a64Hash(buildDatasetFingerprint(input)); + + if (memoKey) { + datasetHashMemo.set(memoKey, { + annotations: input.annotations, + numericAnnotationData: input.numeric_annotation_data, + annotationPredicted: input.annotation_predicted, + hash, + }); + } - return fnv1a64Hash(combined); + return hash; } From 7c68505efe635a09ebcc45755d4a140cb5be4888 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 01:33:07 +0200 Subject: [PATCH 16/31] feat(core): read parquetbundle v3 columnar into typed arrays Prefix-sums the wire's per-row counts into CSR offsets; manifest is validated, not trusted. --- .../src/components/data-loader/data-loader.ts | 12 +- .../data-loader/utils/bundle-v3.test.ts | 483 +++++++++++++ .../components/data-loader/utils/bundle-v3.ts | 636 ++++++++++++++++++ .../data-loader/utils/bundle.test.ts | 30 +- .../components/data-loader/utils/bundle.ts | 164 +++-- .../data-loader/utils/conversion.ts | 46 +- 6 files changed, 1277 insertions(+), 94 deletions(-) create mode 100644 packages/core/src/components/data-loader/utils/bundle-v3.test.ts create mode 100644 packages/core/src/components/data-loader/utils/bundle-v3.ts diff --git a/packages/core/src/components/data-loader/data-loader.ts b/packages/core/src/components/data-loader/data-loader.ts index d0f16d94..ec72ada4 100644 --- a/packages/core/src/components/data-loader/data-loader.ts +++ b/packages/core/src/components/data-loader/data-loader.ts @@ -7,7 +7,7 @@ import { isParquetBundle, type VisualizationData, type BundleSettings } from '@p import { dataLoaderStyles } from './data-loader.styles'; import { createDataErrorEventDetail, type DataErrorEventDetail } from './data-loader.events'; import { readFileOptimized } from './utils/file-io'; -import { extractRowsFromParquetBundle } from './utils/bundle'; +import { decodeParquetBundle } from './utils/bundle'; import { convertParquetToVisualizationDataOptimized } from './utils/conversion'; import { assertValidFileExtension, @@ -219,16 +219,10 @@ export class DataLoader extends LitElement { } catch (workerError) { // Fallback: main-thread decode (worker unsupported / runtime failure). console.warn('Worker decode failed, falling back to main thread:', workerError); - const extraction = await extractRowsFromParquetBundle(arrayBuffer); - validateRowsBasic(extraction.projections); - visualizationData = await convertParquetToVisualizationDataOptimized(extraction); - settings = extraction.settings; + ({ data: visualizationData, settings } = await decodeParquetBundle(arrayBuffer)); } } else { - const extraction = await extractRowsFromParquetBundle(arrayBuffer); - validateRowsBasic(extraction.projections); - visualizationData = await convertParquetToVisualizationDataOptimized(extraction); - settings = extraction.settings; + ({ data: visualizationData, settings } = await decodeParquetBundle(arrayBuffer)); } this.completeStep(); this.dispatchDataLoaded(visualizationData, settings, source, file); diff --git a/packages/core/src/components/data-loader/utils/bundle-v3.test.ts b/packages/core/src/components/data-loader/utils/bundle-v3.test.ts new file mode 100644 index 00000000..41128236 --- /dev/null +++ b/packages/core/src/components/data-loader/utils/bundle-v3.test.ts @@ -0,0 +1,483 @@ +import { describe, it, expect, vi, afterEach } from 'vitest'; +import { readFileSync } from 'node:fs'; +import { parquetWriteBuffer } from 'hyparquet-writer'; +import { + BUNDLE_DELIMITER_BYTES, + concatenateBuffers, + getProteinAnnotationIndices, + getProteinEvidence, + getProteinScores, + isCsrAnnotationData, + NA_VALUE, + type CsrAnnotationData, + type VisualizationData, +} from '@protspace/utils'; +import { decodeParquetBundle } from './bundle'; +import { collectTransferables } from '../decode-transferables'; + +/** + * Format v3 reader tests. + * + * The fixtures are synthesised here rather than produced by the Python encoder, so the + * cases the encoder cannot easily be talked into (a corrupt manifest, a hit count that + * disagrees with its payload) are reachable. Everything they assert was first checked + * against real `encode_v3` output — see the equivalence suite for the producer-written + * side of the contract. + * + * Byte layout notes that the fixtures depend on: + * - lengths on the wire are PER-ROW / PER-HIT COUNTS, prefix-summed by the reader; + * - every part 1/3/6 column is REQUIRED, so hyparquet hands back typed arrays; + * - payload buffers are little-endian, matching every platform the app runs on. + */ + +const enc = new TextEncoder(); +const utf8 = (text: string) => enc.encode(text); +const i32 = (...values: number[]) => new Uint8Array(new Int32Array(values).buffer); +const f32 = (...values: number[]) => new Uint8Array(new Float32Array(values).buffer); + +type Column = { name: string; data: unknown[] | Int32Array | Float64Array | Float32Array }; + +function part(columns: Column[], kv?: Record): Uint8Array { + return new Uint8Array( + parquetWriteBuffer({ + columnData: columns.map((column) => ({ ...column, nullable: false })) as never, + statistics: false, + ...(kv ? { kvMetadata: Object.entries(kv).map(([key, value]) => ({ key, value })) } : {}), + }), + ); +} + +const payloadPart = (payloads: Record): Uint8Array => + part([ + { name: 'name', data: Object.keys(payloads) }, + { name: 'data', data: Object.values(payloads) }, + ]); + +const bundle = (parts: Uint8Array[]): ArrayBuffer => + concatenateBuffers( + parts.map((p) => p.slice().buffer as ArrayBuffer), + BUNDLE_DELIMITER_BYTES, + ); + +// ── the shared fixture ────────────────────────────────────────────────────────── +// +// 8 proteins. `go_bp` is the interesting column: multi-valued, scored, evidenced, and +// with no hits at all on the FIRST row (P1), an INTERIOR one (P4) and the LAST (P8) — +// the three positions the synthetic-NA insertion has to get right. + +const PROTEIN_IDS = ['P1', 'P2', 'P3', 'P4', 'P5', 'P6', 'P7', 'P8']; + +const MANIFEST = { + idColumn: 'protein_id', + columns: { + organism: { kind: 'categorical' }, + go_bp: { kind: 'multi', scores: true, evidence: true }, + keyword: { kind: 'multi' }, + length: { kind: 'numeric', numericType: 'int' }, + score: { kind: 'numeric', numericType: 'float' }, + }, + projections: [ + { name: 'pca2', dimension: 2 }, + { name: 'umap3', dimension: 3 }, + ], +}; + +/** `null` writes no manifest at all; anything else is stamped verbatim. */ +const annotationsPart = (manifest: unknown = MANIFEST) => + part( + [ + { name: 'protein_id', data: PROTEIN_IDS }, + // -1 at P4: the only row with no organism, so a `__NA__` category is appended. + { name: 'organism', data: new Int32Array([0, 1, 2, -1, 0, 1, 2, 3]) }, + { name: 'go_bp__count', data: new Int32Array([0, 2, 1, 0, 3, 1, 2, 0]) }, + { name: 'keyword__count', data: new Int32Array([1, 3, 2, 1, 0, 4, 1, 2]) }, + { name: 'length', data: new Float64Array([100, 200, NaN, 300, 400, 500, 600, 700]) }, + { name: 'score', data: new Float64Array([0.5, 1.5, 2.5, NaN, 4.5, 5.5, 6.5, 7.5]) }, + ], + { + protspace_format_version: '3', + ...(manifest === null ? {} : { protspace_v3_manifest: JSON.stringify(manifest) }), + }, + ); + +const PROJECTIONS_METADATA = part([ + { name: 'projection_name', data: ['pca2', 'umap3'] }, + { name: 'dimensions', data: new Int32Array([2, 3]) }, + { name: 'info_json', data: ['{"note":"flat"}', '{}'] }, +]); + +const PROJECTIONS = part([ + { name: 'pca2__x', data: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8]) }, + { name: 'pca2__y', data: new Float32Array([1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5]) }, + // P7 and P8 are absent from umap3, so the encoder wrote 0.0 for them (matching v2). + { name: 'umap3__x', data: new Float32Array([10, 20, 30, 40, 50, 60, 0, 0]) }, + { name: 'umap3__y', data: new Float32Array([11, 21, 31, 41, 51, 61, 0, 0]) }, + { name: 'umap3__z', data: new Float32Array([0.25, 0.5, 0.75, 1, 1.25, 1.5, 0, 0]) }, +]); + +const PAYLOADS: Record = { + 'dict:organism': utf8('HumanMouseYeastFly'), + 'dict:organism:len': i32(5, 5, 5, 3), + 'dict:go_bp': utf8('bindingapoptosistransport'), + 'dict:go_bp:len': i32(7, 9, 9), + 'csr:go_bp': i32(0, 1, 2, 0, 1, 2, 1, 0, 2), + // Hit 3 is the first hit of P5, immediately after the empty interior row P4: its + // score is what an off-by-one in the inserted-NA `hitEnd` would steal. + 'score_count:go_bp': i32(2, 0, 1, 1, 0, 0, 0, 3, 0), + 'scores:go_bp': f32(1.5, 2.5, 9.75, 4, 0.5, 0.25, 0.125), + 'evidence:go_bp': i32(-1, 0, 1, -1, -1, -1, 0, -1, -1), + 'dict:__evidence': utf8('IDAECO:0000269'), + 'dict:__evidence:len': i32(3, 11), + 'dict:keyword': utf8('alphabetagamma'), + 'dict:keyword:len': i32(5, 4, 5), + 'csr:keyword': i32(0, 0, 1, 2, 1, 2, 2, 0, 1, 2, 0, 1, 2, 0), +}; + +const EMPTY = new Uint8Array(0); + +/** Six parts, with the zero-byte settings and statistics slots the writer emits. */ +const v3Bundle = (overrides: Record = {}) => + bundle( + [annotationsPart(), PROJECTIONS_METADATA, PROJECTIONS, EMPTY, EMPTY, payloadPart(PAYLOADS)].map( + (fallback, index) => overrides[index] ?? fallback, + ), + ); + +const labelsOf = (data: VisualizationData, key: string, protein: number) => + getProteinAnnotationIndices(data.annotation_data[key], protein).map( + (index) => data.annotations[key].values[index], + ); + +describe('parquetbundle format v3', () => { + afterEach(() => vi.restoreAllMocks()); + + it('reads a six-part bundle with zero-byte settings and statistics slots', async () => { + const { data, settings } = await decodeParquetBundle(v3Bundle()); + + expect(settings).toBeNull(); + expect(data.statistics).toBeUndefined(); + expect(data.protein_ids).toEqual(PROTEIN_IDS); + }); + + it('decodes a categorical column and routes its missing row to __NA__', async () => { + const { data } = await decodeParquetBundle(v3Bundle()); + + expect(data.annotations.organism).toEqual({ + kind: 'categorical', + values: ['Human', 'Mouse', 'Yeast', 'Fly', NA_VALUE], + colors: expect.any(Array), + shapes: expect.any(Array), + }); + // Palette assignment must be code-indexed, exactly as the v1/v2 reader does it. + expect(data.annotations.organism.colors).toHaveLength(5); + expect(data.annotations.organism.shapes.every((shape) => shape === 'circle')).toBe(true); + // Plain Int32Array storage for a single-valued column — no CSR, no boxed arrays. + expect(data.annotation_data.organism).toBeInstanceOf(Int32Array); + expect(Array.from(data.annotation_data.organism as Int32Array)).toEqual([ + 0, 1, 2, 4, 0, 1, 2, 3, + ]); + }); + + it('prefix-sums per-row hit counts into CSR offsets', async () => { + const { data } = await decodeParquetBundle(v3Bundle()); + + const csr = data.annotation_data.go_bp as CsrAnnotationData; + expect(isCsrAnnotationData(csr)).toBe(true); + expect(csr.length).toBe(8); + // Counts [0,2,1,0,3,1,2,0] plus one inserted __NA__ hit for each of the three + // empty rows (first, interior, last). + expect(Array.from(csr.end)).toEqual([1, 3, 4, 5, 8, 9, 11, 12]); + expect(Array.from(csr.codes)).toEqual([3, 0, 1, 2, 3, 0, 1, 2, 1, 0, 2, 3]); + }); + + it('gives empty rows at the first, interior and last positions the __NA__ category', async () => { + const { data } = await decodeParquetBundle(v3Bundle()); + + expect(data.annotations.go_bp.values).toEqual(['binding', 'apoptosis', 'transport', NA_VALUE]); + expect(labelsOf(data, 'go_bp', 0)).toEqual([NA_VALUE]); + expect(labelsOf(data, 'go_bp', 3)).toEqual([NA_VALUE]); + expect(labelsOf(data, 'go_bp', 7)).toEqual([NA_VALUE]); + expect(labelsOf(data, 'go_bp', 1)).toEqual(['binding', 'apoptosis']); + expect(labelsOf(data, 'go_bp', 4)).toEqual(['binding', 'apoptosis', 'transport']); + }); + + it('keeps scores aligned with their hits across the inserted __NA__ hits', async () => { + const { data } = await decodeParquetBundle(v3Bundle()); + + expect(getProteinScores(data, 0, 'go_bp')).toEqual([null]); + expect(getProteinScores(data, 1, 'go_bp')).toEqual([[1.5, 2.5], null]); + expect(getProteinScores(data, 2, 'go_bp')).toEqual([[9.75]]); + // P4 is empty and P5's first hit is scored: the inserted __NA__ hit must own no + // score, and P5's must keep the one it was written with. + expect(getProteinScores(data, 3, 'go_bp')).toEqual([null]); + expect(getProteinScores(data, 4, 'go_bp')).toEqual([[4], null, null]); + expect(getProteinScores(data, 5, 'go_bp')).toEqual([null]); + expect(getProteinScores(data, 6, 'go_bp')).toEqual([[0.5, 0.25, 0.125], null]); + expect(getProteinScores(data, 7, 'go_bp')).toEqual([null]); + }); + + it('keeps evidence aligned with its hits and resolves the global evidence dictionary', async () => { + const { data } = await decodeParquetBundle(v3Bundle()); + + expect(getProteinEvidence(data, 1, 'go_bp')).toEqual([null, 'IDA']); + expect(getProteinEvidence(data, 2, 'go_bp')).toEqual(['ECO:0000269']); + expect(getProteinEvidence(data, 5, 'go_bp')).toEqual(['IDA']); + expect(getProteinEvidence(data, 7, 'go_bp')).toEqual([null]); + }); + + it('leaves a multi column with neither scores nor evidence without those payloads', async () => { + const { data } = await decodeParquetBundle(v3Bundle()); + + expect(labelsOf(data, 'keyword', 1)).toEqual(['alpha', 'beta', 'gamma']); + expect(labelsOf(data, 'keyword', 4)).toEqual([NA_VALUE]); + expect(data.annotation_scores_csr?.keyword).toBeUndefined(); + expect(data.annotation_evidence_csr?.keyword).toBeUndefined(); + expect(getProteinScores(data, 1, 'keyword')).toEqual([]); + }); + + it('reads numeric columns from float64 with NaN meaning missing', async () => { + const { data } = await decodeParquetBundle(v3Bundle()); + + expect(data.annotations.length).toMatchObject({ kind: 'numeric', numericType: 'int' }); + expect(data.annotations.score).toMatchObject({ kind: 'numeric', numericType: 'float' }); + expect(data.numeric_annotation_data?.length).toEqual([100, 200, null, 300, 400, 500, 600, 700]); + expect(data.numeric_annotation_data?.score).toEqual([0.5, 1.5, 2.5, null, 4.5, 5.5, 6.5, 7.5]); + // The manifest is authoritative: an int32 code column must never be read as numeric. + expect(data.numeric_annotation_data?.organism).toBeUndefined(); + expect(data.numeric_annotation_data?.go_bp).toBeUndefined(); + }); + + it('interleaves the wide axis columns into 2D and 3D projections', async () => { + const { data } = await decodeParquetBundle(v3Bundle()); + + const [pca2, umap3] = data.projections; + expect(pca2).toMatchObject({ name: 'pca2', dimension: 2 }); + expect(Array.from(pca2.data)).toEqual([ + 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5, 6, 6.5, 7, 7.5, 8, 8.5, + ]); + expect(pca2.metadata).toMatchObject({ dimension: 2, dimensions: 2, note: 'flat' }); + + expect(umap3).toMatchObject({ name: 'umap3', dimension: 3 }); + expect(Array.from(umap3.data.slice(0, 6))).toEqual([10, 11, 0.25, 20, 21, 0.5]); + // A protein absent from a projection sits at the origin, matching v2. + expect(Array.from(umap3.data.slice(18))).toEqual([0, 0, 0, 0, 0, 0]); + }); + + it('takes the typed-array fast path for every column', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}); + await decodeParquetBundle(v3Bundle()); + expect(warn).not.toHaveBeenCalled(); + }); + + it('parses the settings and statistics parts when they are present', async () => { + const settingsPart = part([ + { + name: 'settings_json', + data: [JSON.stringify({ legendSettings: {}, exportOptions: {} })], + }, + ]); + const statisticsPart = new Uint8Array( + readFileSync(new URL('./__fixtures__/stats-sample-statistics.parquet', import.meta.url)), + ); + + const { data, settings } = await decodeParquetBundle( + v3Bundle({ 3: settingsPart, 4: statisticsPart }), + ); + + expect(settings).toEqual({ legendSettings: {}, exportOptions: {} }); + expect(new Uint8Array(data.statistics!)).toEqual(statisticsPart); + expect(data.statisticsRows!.length).toBeGreaterThan(0); + }); + + describe('rejects a bundle whose manifest cannot be trusted', () => { + const cases: [string, unknown, RegExp][] = [ + ['no manifest at all', null, /carries no "protspace_v3_manifest"/], + [ + 'an unknown column kind', + { ...MANIFEST, columns: { organism: { kind: 'blob' } } }, + /unknown kind "blob"/, + ], + [ + 'a column that part 1 does not have', + { ...MANIFEST, columns: { ...MANIFEST.columns, ghost: { kind: 'categorical' } } }, + /declares column "ghost" but part 1 has no "ghost"/, + ], + [ + 'a multi column named as if it were single-valued', + { ...MANIFEST, columns: { ...MANIFEST.columns, organism: { kind: 'multi' } } }, + /part 1 has no "organism__count"/, + ], + [ + 'an id column that is not in the schema', + { ...MANIFEST, idColumn: 'accession' }, + /idColumn "accession" is not a column of part 1/, + ], + [ + 'a projection dimension other than 2 or 3', + { ...MANIFEST, projections: [{ name: 'pca2', dimension: 4 }] }, + /dimension 4, expected 2 or 3/, + ], + [ + 'a projection column part 3 does not have', + { ...MANIFEST, projections: [{ name: 'nope', dimension: 2 }] }, + /part 3 has no nope__x/, + ], + [ + 'an unknown numericType', + { + ...MANIFEST, + columns: { ...MANIFEST.columns, length: { kind: 'numeric', numericType: 'i8' } }, + }, + /unknown numericType "i8"/, + ], + ]; + + for (const [label, manifest, message] of cases) { + it(label, async () => { + await expect( + decodeParquetBundle(v3Bundle({ 0: annotationsPart(manifest) })), + ).rejects.toThrow(message); + }); + } + + it('a manifest that is not JSON', async () => { + const broken = part([{ name: 'protein_id', data: PROTEIN_IDS }], { + protspace_format_version: '3', + protspace_v3_manifest: '{not json', + }); + await expect(decodeParquetBundle(v3Bundle({ 0: broken }))).rejects.toThrow( + /manifest is not valid JSON/, + ); + }); + }); + + describe('rejects payloads that disagree with part 1', () => { + it('hit counts that do not sum to the CSR code count', async () => { + const payloads = { ...PAYLOADS, 'csr:go_bp': i32(0, 1, 2) }; + await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( + /hit counts sum to 9 but csr:go_bp holds 3 codes/, + ); + }); + + it('a code outside the column dictionary', async () => { + const payloads = { ...PAYLOADS, 'csr:go_bp': i32(0, 1, 2, 0, 1, 2, 1, 0, 7) }; + await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( + /hit 8 has code 7, outside its 3 labels/, + ); + }); + + it('a categorical code outside the column dictionary', async () => { + const payloads = { + ...PAYLOADS, + 'dict:organism:len': i32(5, 5, 5), + 'dict:organism': utf8('HumanMouseYeast'), + }; + await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( + /row 7 has code 3, outside its 3 labels/, + ); + }); + + it('score counts that do not sum to the score count', async () => { + const payloads = { ...PAYLOADS, 'scores:go_bp': f32(1.5, 2.5) }; + await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( + /score counts sum to 7 but scores:go_bp holds 2/, + ); + }); + + it('a dictionary blob shorter than its label lengths', async () => { + const payloads = { ...PAYLOADS, 'dict:organism': utf8('HumanMouse') }; + await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( + /declares a label past the end of its blob/, + ); + }); + + it('a payload whose byte length is not a multiple of 4', async () => { + const payloads = { ...PAYLOADS, 'csr:go_bp': utf8('xyz') }; + await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( + /is 3 bytes, not a multiple of 4/, + ); + }); + + it('a missing payloads part', async () => { + await expect(decodeParquetBundle(v3Bundle({ 5: EMPTY }))).rejects.toThrow( + /carries no payloads part/, + ); + }); + }); + + it('decodes non-ASCII labels by byte range, not character offset', async () => { + // 'Mü' is three bytes but two characters, so slicing the decoded blob by byte + // offsets would shear every later label. + const payloads = { + ...PAYLOADS, + 'dict:organism': utf8('HumanMüYeastFly'), + 'dict:organism:len': i32(5, 3, 5, 3), + }; + const { data } = await decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) })); + expect(data.annotations.organism.values).toEqual(['Human', 'Mü', 'Yeast', 'Fly', NA_VALUE]); + }); + + it('still reads a bundle whose columns were written nullable, and says so once', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const nullable = new Uint8Array( + parquetWriteBuffer({ + columnData: [ + { name: 'protein_id', data: PROTEIN_IDS, nullable: false }, + { name: 'organism', data: [0, 1, 2, -1, 0, 1, 2, 3], type: 'INT32', nullable: true }, + { name: 'go_bp__count', data: new Int32Array([0, 2, 1, 0, 3, 1, 2, 0]), nullable: false }, + { + name: 'keyword__count', + data: new Int32Array([1, 3, 2, 1, 0, 4, 1, 2]), + nullable: false, + }, + { + name: 'length', + data: [100, 200, null, 300, 400, 500, 600, 700], + type: 'DOUBLE', + nullable: true, + }, + { + name: 'score', + data: new Float64Array([0.5, 1.5, 2.5, NaN, 4.5, 5.5, 6.5, 7.5]), + nullable: false, + }, + ] as never, + statistics: false, + kvMetadata: [ + { key: 'protspace_format_version', value: '3' }, + { key: 'protspace_v3_manifest', value: JSON.stringify(MANIFEST) }, + ], + }), + ); + + const { data } = await decodeParquetBundle(v3Bundle({ 0: nullable })); + + expect(Array.from(data.annotation_data.organism as Int32Array)).toEqual([ + 0, 1, 2, 4, 0, 1, 2, 3, + ]); + // A null in a column the manifest calls numeric reads as missing, not as 0. + expect(data.numeric_annotation_data?.length).toEqual([100, 200, null, 300, 400, 500, 600, 700]); + expect(warn).toHaveBeenCalledTimes(1); + expect(warn.mock.calls[0][0]).toMatch(/did not decode to a typed array/); + }); + + it('collects every bulk buffer exactly once and they all transfer', async () => { + const { data } = await decodeParquetBundle(v3Bundle()); + const transfer = collectTransferables(data); + + expect(new Set(transfer).size).toBe(transfer.length); + // 2 projections + organism codes + 2 CSR (end + codes) + scores (hitEnd + values) + // + evidence codes. + expect(transfer).toHaveLength(10); + + const sources = [ + ...data.projections.map((projection) => projection.data), + data.annotation_data.organism as Int32Array, + (data.annotation_data.go_bp as CsrAnnotationData).codes, + data.annotation_scores_csr!.go_bp.values, + data.annotation_evidence_csr!.go_bp.codes, + ]; + structuredClone(data, { transfer }); + expect(sources.every((array) => array.byteLength === 0)).toBe(true); + }); +}); diff --git a/packages/core/src/components/data-loader/utils/bundle-v3.ts b/packages/core/src/components/data-loader/utils/bundle-v3.ts new file mode 100644 index 00000000..074e3d9b --- /dev/null +++ b/packages/core/src/components/data-loader/utils/bundle-v3.ts @@ -0,0 +1,636 @@ +/** + * Reader for `.parquetbundle` format v3 — the columnar annotation encoding written by + * `apps/protspace/src/protspace/data/io/bundle_v3.py`, which is the specification for + * everything below. + * + * v1 and v2 stringify every annotation cell, so loading them means one JS object per + * row plus a re-split and re-dictionary-coding of every string. v3 does that work at + * write time: part 1 carries int32 dictionary codes (or per-row CSR hit counts) and + * float64 numerics, part 3 carries wide float32 projections, and part 6 carries the + * label dictionaries and CSR code/score/evidence payloads as raw little-endian buffers. + * This reader therefore never parses a string that is not a label, and hands the worker + * typed arrays it can transfer instead of structured-clone. + * + * Two wire details drive most of the code here: + * + * - **Lengths are per-element counts, never cumulative offsets.** Offsets are + * near-incompressible; their first differences are not. Every `__count`, + * `score_count:` and `dict::len` family is prefix-summed here into the + * cumulative offsets the in-memory `CsrAnnotationData` / `CsrScores` types use. + * - **Every part 1/3/6 column is REQUIRED and PLAIN**, which is the only shape + * hyparquet decodes straight into a typed array. A column that arrives as a plain + * array still reads correctly (see the fallback in {@link writeChunk}) but about 4x + * slower, and is a bug on the writer side, so it is logged. + */ + +import { parquetMetadata, parquetRead, parquetReadObjects, type FileMetaData } from 'hyparquet'; +import { + NA_DEFAULT_COLOR, + NA_VALUE, + type Annotation, + type AnnotationData, + type BundleSettings, + type CsrEvidence, + type CsrScores, + type Projection, + type VisualizationData, +} from '@protspace/utils'; +import { assertValidParquetMagic } from './validation'; +import { extractSettings, extractStatistics } from './bundle'; +import { + appendSyntheticNACategoryToCodes, + buildProjectionsMetadataMap, + carryStatistics, + createNumericAnnotation, + generateColorsAndShapes, + normalizeEatCompanionColumns, +} from './conversion'; +import type { Rows } from './types'; + +/** Key-value metadata key part 1 carries the v3 manifest under. */ +const MANIFEST_KEY = 'protspace_v3_manifest'; + +/** Payload name of the dictionary every column's evidence codes index into. */ +const EVIDENCE_DICT_NAME = '__evidence'; + +const AXES = ['x', 'y', 'z'] as const; + +const DECODER = new TextDecoder(); + +type V3ColumnKind = 'categorical' | 'multi' | 'numeric'; + +interface V3ColumnManifest { + kind: V3ColumnKind; + /** Only meaningful for `kind: 'numeric'`; defaults to float when absent. */ + numericType?: 'int' | 'float'; + /** Only meaningful for `kind: 'multi'`: a `scores:` payload exists. */ + scores?: boolean; + /** Only meaningful for `kind: 'multi'`: an `evidence:` payload exists. */ + evidence?: boolean; +} + +interface V3Manifest { + idColumn: string; + columns: Record; + projections: { name: string; dimension: 2 | 3 }[]; +} + +/** Physical part-1 column backing a manifest column: multi stores per-row hit counts. */ +function physicalColumn(name: string, kind: V3ColumnKind): string { + return kind === 'multi' ? `${name}__count` : name; +} + +/** Leaf (data) column names of a parquet schema; the root element carries no type. */ +function leafColumnNames(metadata: FileMetaData): Set { + const names = new Set(); + for (const field of metadata.schema) { + if (field.name && field.type) names.add(field.name); + } + return names; +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +/** + * Parse and **validate** the v3 manifest against part 1's own schema. + * + * This is a trust boundary: the manifest is the only thing that says how the int32 + * columns below should be interpreted, and a wrong `kind` would silently turn a code + * column into a numeric annotation (or index a dictionary that isn't there). Every + * mismatch throws with the offending name rather than being repaired, so a broken + * producer is reported instead of half-rendered. + */ +function readManifest(metadata: FileMetaData): V3Manifest { + const raw = metadata.key_value_metadata?.find((entry) => entry.key === MANIFEST_KEY)?.value; + if (!raw) { + throw new Error(`Bundle declares format v3 but carries no "${MANIFEST_KEY}" metadata`); + } + + let parsed: unknown; + try { + parsed = JSON.parse(raw); + } catch (error) { + throw new Error(`v3 manifest is not valid JSON: ${(error as Error).message}`); + } + if (!isRecord(parsed)) throw new Error('v3 manifest is not a JSON object'); + + const schemaColumns = leafColumnNames(metadata); + + const { idColumn, columns, projections } = parsed; + if (typeof idColumn !== 'string' || !schemaColumns.has(idColumn)) { + throw new Error(`v3 manifest idColumn "${String(idColumn)}" is not a column of part 1`); + } + if (!isRecord(columns)) throw new Error('v3 manifest has no "columns" object'); + if (!Array.isArray(projections)) throw new Error('v3 manifest has no "projections" array'); + + const validated: Record = {}; + for (const [name, entry] of Object.entries(columns)) { + if (!isRecord(entry)) + throw new Error(`v3 manifest entry for column "${name}" is not an object`); + const { kind, numericType } = entry; + if (kind !== 'categorical' && kind !== 'multi' && kind !== 'numeric') { + throw new Error(`v3 manifest column "${name}" has unknown kind "${String(kind)}"`); + } + if (numericType != null && numericType !== 'int' && numericType !== 'float') { + throw new Error( + `v3 manifest column "${name}" has unknown numericType "${String(numericType)}"`, + ); + } + const physical = physicalColumn(name, kind); + if (!schemaColumns.has(physical)) { + throw new Error(`v3 manifest declares column "${name}" but part 1 has no "${physical}"`); + } + validated[name] = { + kind, + ...(numericType != null ? { numericType } : {}), + ...(entry.scores === true ? { scores: true } : {}), + ...(entry.evidence === true ? { evidence: true } : {}), + }; + } + + const validatedProjections: V3Manifest['projections'] = []; + const seen = new Set(); + for (const entry of projections) { + if (!isRecord(entry)) throw new Error('v3 manifest projection entry is not an object'); + const { name, dimension } = entry; + if (typeof name !== 'string' || !name) { + throw new Error(`v3 manifest projection has an invalid name "${String(name)}"`); + } + if (seen.has(name)) throw new Error(`v3 manifest declares projection "${name}" twice`); + seen.add(name); + if (dimension !== 2 && dimension !== 3) { + throw new Error( + `v3 projection "${name}" has dimension ${String(dimension)}, expected 2 or 3`, + ); + } + validatedProjections.push({ name, dimension }); + } + + return { idColumn, columns: validated, projections: validatedProjections }; +} + +type ColumnTarget = Int32Array | Float64Array | string[]; + +/** + * Copy one decoded chunk into its preallocated column at `rowStart`. + * + * The fast path is the whole point of v3: a REQUIRED PLAIN column arrives as a typed + * array and lands with a single `set`. Anything else — a column the producer wrote + * nullable or dictionary-encoded — still decodes correctly through the element loop, + * which is why `onPlainArray` reports rather than throws. + */ +function writeChunk( + target: ColumnTarget, + columnName: string, + columnData: ArrayLike, + rowStart: number, + onPlainArray: (columnName: string) => void, +): void { + if (Array.isArray(target)) { + for (let i = 0; i < columnData.length; i++) { + const value = columnData[i]; + target[rowStart + i] = typeof value === 'string' ? value : String(value ?? ''); + } + return; + } + + if ( + columnData instanceof Int32Array || + columnData instanceof Float64Array || + columnData instanceof Float32Array + ) { + target.set(columnData, rowStart); + return; + } + + onPlainArray(columnName); + const missing = target instanceof Int32Array ? -1 : NaN; + for (let i = 0; i < columnData.length; i++) { + const value = columnData[i]; + target[rowStart + i] = value == null ? missing : Number(value); + } +} + +/** Preallocate one array per declared column and fill it chunk by chunk. */ +async function readAnnotationColumns( + part: ArrayBuffer, + metadata: FileMetaData, + manifest: V3Manifest, + numRows: number, +): Promise> { + const targets = new Map(); + targets.set(manifest.idColumn, new Array(numRows).fill('')); + for (const [name, column] of Object.entries(manifest.columns)) { + targets.set( + physicalColumn(name, column.kind), + column.kind === 'numeric' ? new Float64Array(numRows) : new Int32Array(numRows), + ); + } + + let warned = false; + const onPlainArray = (columnName: string) => { + if (warned) return; + warned = true; + console.warn( + `v3 bundle column "${columnName}" did not decode to a typed array — it was probably ` + + 'written nullable or dictionary-encoded. The bundle still loads, about 4x slower; ' + + 'fix the writer (every v3 column must be REQUIRED and PLAIN).', + ); + }; + + await parquetRead({ + file: part, + metadata, + columns: [...targets.keys()], + onChunk: ({ columnName, columnData, rowStart }) => { + const target = targets.get(columnName); + if (target) writeChunk(target, columnName, columnData, rowStart, onPlainArray); + }, + }); + + return targets; +} + +/** + * Read part 3 into one flat `Float32Array(N * dimension)` per projection. + * + * The wire is one column per axis (`__x`, `__y`, `__z`), so the interleave into + * the renderer's stride-major layout happens right in the chunk callback: no + * per-projection intermediate and no second pass. A protein absent from a projection + * keeps the zero the allocation gave it, which is what v2 produced too. + */ +async function readProjections( + part: ArrayBuffer, + manifest: V3Manifest, + numRows: number, + metadataMap: ReadonlyMap>, +): Promise { + assertValidParquetMagic(part); + const metadata = parquetMetadata(part); + const schemaColumns = leafColumnNames(metadata); + + const axisTargets = new Map(); + const projections: Projection[] = []; + + for (const { name, dimension } of manifest.projections) { + const data = new Float32Array(numRows * dimension); + for (let axis = 0; axis < dimension; axis++) { + const column = `${name}__${AXES[axis]}`; + if (!schemaColumns.has(column)) { + throw new Error( + `v3 projection "${name}" declares ${dimension}D but part 3 has no ${column}`, + ); + } + axisTargets.set(column, { data, dimension, axis }); + } + projections.push({ + name, + data, + dimension, + metadata: { ...(metadataMap.get(name) ?? {}), dimension }, + }); + } + + if (axisTargets.size > 0) { + await parquetRead({ + file: part, + metadata, + columns: [...axisTargets.keys()], + onChunk: ({ columnName, columnData, rowStart }) => { + const target = axisTargets.get(columnName); + if (!target) return; + const { data, dimension, axis } = target; + for (let i = 0; i < columnData.length; i++) { + data[(rowStart + i) * dimension + axis] = columnData[i] as number; + } + }, + }); + } + + return projections; +} + +/** Part 6 as `name -> raw little-endian bytes`. */ +async function readPayloads(part: ArrayBuffer): Promise> { + assertValidParquetMagic(part); + // utf8: false keeps the `data` column as raw bytes. The `name` column carries a + // STRING logical type, which hyparquet decodes regardless of this flag. + const rows = await parquetReadObjects({ file: part, utf8: false }); + const payloads = new Map(); + for (const row of rows) { + const name = typeof row.name === 'string' ? row.name : DECODER.decode(row.name as Uint8Array); + payloads.set(name, row.data as Uint8Array); + } + return payloads; +} + +/** + * A payload as an aligned typed array. + * + * hyparquet hands back a `Uint8Array` **view into the page buffer**, at an arbitrary + * byte offset — so it is copied rather than wrapped. Wrapping would both risk an + * alignment error and pin (or, if transferred, detach) the whole decoded page. + */ +function asTypedPayload( + payloads: ReadonlyMap, + name: string, + ctor: new (buffer: ArrayBuffer) => T, +): T { + const bytes = payloads.get(name); + if (!bytes) throw new Error(`v3 bundle is missing the "${name}" payload`); + if (bytes.byteLength % 4 !== 0) { + throw new Error(`v3 payload "${name}" is ${bytes.byteLength} bytes, not a multiple of 4`); + } + return new ctor(bytes.slice().buffer); +} + +/** + * Labels of one dictionary payload, in code order. + * + * The blob is the utf8 concatenation and `:len` holds each label's **byte** length. + * Decoding the blob once and slicing it is only valid while character offsets equal + * byte offsets, i.e. while the blob is pure ASCII — which covers most columns; the + * moment it is not, each label is decoded from its own byte range instead. + */ +function readLabels(payloads: ReadonlyMap, name: string): string[] { + const lengths = asTypedPayload(payloads, `dict:${name}:len`, Int32Array); + const bytes = payloads.get(`dict:${name}`); + if (!bytes) throw new Error(`v3 bundle is missing the "dict:${name}" payload`); + + const labels = new Array(lengths.length); + const text = DECODER.decode(bytes); + const ascii = text.length === bytes.byteLength; + let at = 0; + for (let i = 0; i < lengths.length; i++) { + const length = lengths[i]; + if (length < 0 || at + length > bytes.byteLength) { + throw new Error(`v3 dictionary "${name}" declares a label past the end of its blob`); + } + labels[i] = ascii + ? text.slice(at, at + length) + : DECODER.decode(bytes.subarray(at, at + length)); + at += length; + } + if (at !== bytes.byteLength) { + throw new Error( + `v3 dictionary "${name}" label lengths cover ${at} of ${bytes.byteLength} blob bytes`, + ); + } + return labels; +} + +/** Per-element counts to the cumulative offsets the in-memory CSR types use. */ +function prefixSum(counts: Int32Array, what: string): Int32Array { + const end = new Int32Array(counts.length); + let running = 0; + for (let i = 0; i < counts.length; i++) { + const count = counts[i]; + if (count < 0) throw new Error(`v3 ${what} has a negative count (${count}) at index ${i}`); + running += count; + end[i] = running; + } + return end; +} + +interface CsrColumn { + end: Int32Array; + codes: Int32Array; + scores: CsrScores | null; + evidence: CsrEvidence | null; +} + +/** Assemble one multi-valued column's CSR storage plus its score/evidence payloads. */ +function readCsrColumn( + name: string, + column: V3ColumnManifest, + counts: Int32Array, + labelCount: number, + payloads: ReadonlyMap, + evidenceDict: () => readonly string[], +): CsrColumn { + const codes = asTypedPayload(payloads, `csr:${name}`, Int32Array); + const end = prefixSum(counts, `column "${name}" hit counts`); + const total = counts.length > 0 ? end[counts.length - 1] : 0; + if (total !== codes.length) { + throw new Error( + `v3 column "${name}" hit counts sum to ${total} but csr:${name} holds ${codes.length} codes`, + ); + } + for (let hit = 0; hit < codes.length; hit++) { + if (codes[hit] < 0 || codes[hit] >= labelCount) { + throw new Error( + `v3 column "${name}" hit ${hit} has code ${codes[hit]}, outside its ${labelCount} labels`, + ); + } + } + + let scores: CsrScores | null = null; + if (column.scores) { + const scoreCounts = asTypedPayload(payloads, `score_count:${name}`, Int32Array); + if (scoreCounts.length !== codes.length) { + throw new Error( + `v3 column "${name}" has ${scoreCounts.length} score counts for ${codes.length} hits`, + ); + } + const values = asTypedPayload(payloads, `scores:${name}`, Float32Array); + const hitEnd = prefixSum(scoreCounts, `column "${name}" score counts`); + const scoreTotal = hitEnd.length > 0 ? hitEnd[hitEnd.length - 1] : 0; + if (scoreTotal !== values.length) { + throw new Error( + `v3 column "${name}" score counts sum to ${scoreTotal} but scores:${name} holds ${values.length}`, + ); + } + scores = { hitEnd, values }; + } + + let evidence: CsrEvidence | null = null; + if (column.evidence) { + const evidenceCodes = asTypedPayload(payloads, `evidence:${name}`, Int32Array); + if (evidenceCodes.length !== codes.length) { + throw new Error( + `v3 column "${name}" has ${evidenceCodes.length} evidence codes for ${codes.length} hits`, + ); + } + evidence = { codes: evidenceCodes, dict: evidenceDict() }; + } + + return { end, codes, scores, evidence }; +} + +/** + * Route rows with no hits at all to a synthetic `__NA__` category, the way + * `appendSyntheticNACategory` does for the nested storage shape. + * + * CSR needs a rebuild rather than an in-place patch: an empty row owns no hit slot to + * write the category into. One lockstep pass therefore inserts a hit per empty row, + * carrying `-1` into evidence and a repeated running total into `hitEnd` — the inserted + * hit contributes no score, so every original hit keeps the exact cumulative it had. + */ +function insertNAForEmptyRows( + csr: CsrColumn, + labels: string[], + colors: string[], + shapes: string[], +): CsrColumn { + const numRows = csr.end.length; + let empty = 0; + for (let i = 0; i < numRows; i++) { + if ((i === 0 ? 0 : csr.end[i - 1]) === csr.end[i]) empty++; + } + if (empty === 0) return csr; + + const naIndex = labels.length; + labels.push(NA_VALUE); + colors.push(NA_DEFAULT_COLOR); + shapes.push('circle'); + + const total = csr.codes.length + empty; + const codes = new Int32Array(total); + const end = new Int32Array(numRows); + const evidenceCodes = csr.evidence ? new Int32Array(total) : null; + const hitEnd = csr.scores ? new Int32Array(total) : null; + + let write = 0; + for (let i = 0; i < numRows; i++) { + const from = i === 0 ? 0 : csr.end[i - 1]; + const to = csr.end[i]; + if (from === to) { + codes[write] = naIndex; + if (evidenceCodes) evidenceCodes[write] = -1; + if (hitEnd) hitEnd[write] = from === 0 ? 0 : csr.scores!.hitEnd[from - 1]; + write++; + } else { + for (let hit = from; hit < to; hit++) { + codes[write] = csr.codes[hit]; + if (evidenceCodes) evidenceCodes[write] = csr.evidence!.codes[hit]; + if (hitEnd) hitEnd[write] = csr.scores!.hitEnd[hit]; + write++; + } + } + end[i] = write; + } + + return { + end, + codes, + scores: csr.scores ? { hitEnd: hitEnd!, values: csr.scores.values } : null, + evidence: csr.evidence ? { codes: evidenceCodes!, dict: csr.evidence.dict } : null, + }; +} + +/** + * Read a format v3 bundle into `VisualizationData`. + * + * `parts` comes from `splitBundleParts`; `metadata` is part 1's already-parsed footer. + */ +export async function readV3Bundle( + parts: readonly (ArrayBuffer | null)[], + metadata: FileMetaData, +): Promise<{ data: VisualizationData; settings: BundleSettings | null }> { + const [part1, part2, part3, part4, part5, part6] = parts; + if (!part1 || !part2 || !part3) { + throw new Error('Parquetbundle is missing one of its three required core parts'); + } + if (!part6) { + throw new Error('Bundle declares format v3 but carries no payloads part (part 6)'); + } + + const manifest = readManifest(metadata); + const numRows = Number(metadata.num_rows); + + const columns = await readAnnotationColumns(part1, metadata, manifest, numRows); + const protein_ids = columns.get(manifest.idColumn) as string[]; + + assertValidParquetMagic(part2); + const projectionsMetadata = (await parquetReadObjects({ file: part2 })) as Rows; + const projections = await readProjections( + part3, + manifest, + numRows, + buildProjectionsMetadataMap(projectionsMetadata), + ); + + const payloads = await readPayloads(part6); + let evidenceDict: readonly string[] | null = null; + const readEvidenceDict = (): readonly string[] => + (evidenceDict ??= readLabels(payloads, EVIDENCE_DICT_NAME)); + + const annotations: Record = {}; + const annotation_data: Record = {}; + const numeric_annotation_data: Record = {}; + const annotation_scores_csr: Record = {}; + const annotation_evidence_csr: Record = {}; + + for (const [name, column] of Object.entries(manifest.columns)) { + const stored = columns.get(physicalColumn(name, column.kind))!; + + if (column.kind === 'numeric') { + const raw = stored as Float64Array; + const values = new Array(numRows); + for (let i = 0; i < numRows; i++) values[i] = Number.isFinite(raw[i]) ? raw[i] : null; + numeric_annotation_data[name] = values; + annotations[name] = createNumericAnnotation(column.numericType ?? 'float'); + continue; + } + + const labels = readLabels(payloads, name); + const { colors, shapes } = generateColorsAndShapes('kellys', labels.length); + + if (column.kind === 'categorical') { + const codes = stored as Int32Array; + for (let i = 0; i < numRows; i++) { + if (codes[i] >= labels.length || codes[i] < -1) { + throw new Error( + `v3 column "${name}" row ${i} has code ${codes[i]}, outside its ${labels.length} labels`, + ); + } + } + appendSyntheticNACategoryToCodes(labels, colors, shapes, codes); + annotation_data[name] = codes; + } else { + const csr = insertNAForEmptyRows( + readCsrColumn( + name, + column, + stored as Int32Array, + labels.length, + payloads, + readEvidenceDict, + ), + labels, + colors, + shapes, + ); + annotation_data[name] = { kind: 'csr', end: csr.end, codes: csr.codes, length: numRows }; + if (csr.scores) annotation_scores_csr[name] = csr.scores; + if (csr.evidence) annotation_evidence_csr[name] = csr.evidence; + } + + annotations[name] = { kind: 'categorical', values: labels, colors, shapes }; + } + + const data: VisualizationData = { + protein_ids, + projections, + annotations, + annotation_data, + numeric_annotation_data, + annotation_scores: {}, + annotation_evidence: {}, + ...(Object.keys(annotation_scores_csr).length > 0 ? { annotation_scores_csr } : {}), + ...(Object.keys(annotation_evidence_csr).length > 0 ? { annotation_evidence_csr } : {}), + }; + + // Deliberately NOT restoreDeclaredNumericAnnotations: it reads physical parquet types, + // which in v3 would declare every int32 dictionary-code column numeric. The manifest + // is the authority on kind here, and it has already been applied above. + return { + data: carryStatistics(normalizeEatCompanionColumns(data), { + statistics: part5, + statisticsRows: part5 ? await extractStatistics(part5) : null, + }), + settings: part4 ? await extractSettings(part4) : null, + }; +} diff --git a/packages/core/src/components/data-loader/utils/bundle.test.ts b/packages/core/src/components/data-loader/utils/bundle.test.ts index faee81c4..d9ad5b0e 100644 --- a/packages/core/src/components/data-loader/utils/bundle.test.ts +++ b/packages/core/src/components/data-loader/utils/bundle.test.ts @@ -112,21 +112,31 @@ describe('bundle utilities', () => { const bundle = createMockBundle(2); await expect(extractRowsFromParquetBundle(bundle)).rejects.toThrow( - /Expected 2 to 4 delimiters/, + /Expected 2 to 5 delimiters/, ); }); - // 5 parts (settings + statistics) is a layout the Python producer writes, so it - // must pass this gate. It is not asserted here: with mock parts the call still - // rejects during decode, so any assertion would be about the decode error, not - // about acceptance. `tests/contract/bundle.contract.test.ts` proves acceptance - // against a real producer-written 5-part bundle instead. + // 5 parts (settings + statistics) and 6 (format v3, which appends the payloads + // part) are layouts the Python producer writes, so they must pass this gate. The + // 5-part case is not asserted here: with mock parts the call still rejects during + // decode, so any assertion would be about the decode error, not about acceptance. + // `tests/contract/bundle.contract.test.ts` proves acceptance against a real + // producer-written 5-part bundle instead. - it('should reject bundle with 5 delimiters (6 parts)', async () => { - const bundle = createMockBundle(6); + it('should let a 6-part bundle past the gate and fail on its contents instead', async () => { + const error: unknown = await extractRowsFromParquetBundle(createMockBundle(6)).catch( + (reason: unknown) => reason, + ); + + expect(error).toBeInstanceOf(Error); + expect((error as Error).message).not.toMatch(/Expected 2 to 5 delimiters/); + }); + + it('should reject bundle with 6 delimiters (7 parts)', async () => { + const bundle = createMockBundle(7); await expect(extractRowsFromParquetBundle(bundle)).rejects.toThrow( - /Expected 2 to 4 delimiters/, + /Expected 2 to 5 delimiters/, ); }); @@ -134,7 +144,7 @@ describe('bundle utilities', () => { const buffer = createMockParquetBuffer('no delimiter'); await expect(extractRowsFromParquetBundle(buffer)).rejects.toThrow( - /Expected 2 to 4 delimiters/, + /Expected 2 to 5 delimiters/, ); }); }); diff --git a/packages/core/src/components/data-loader/utils/bundle.ts b/packages/core/src/components/data-loader/utils/bundle.ts index 88a85dd5..57ae9e67 100644 --- a/packages/core/src/components/data-loader/utils/bundle.ts +++ b/packages/core/src/components/data-loader/utils/bundle.ts @@ -6,9 +6,12 @@ import { normalizeBundleSettings, type BundleSettings, type ProjectionStatisticRow, + type VisualizationData, } from '@protspace/utils'; import type { Rows, GenericRow } from './types'; -import { assertValidParquetMagic, validateProjectionRows } from './validation'; +import { assertValidParquetMagic, validateProjectionRows, validateRowsBasic } from './validation'; +import { convertParquetToVisualizationDataOptimized } from './conversion'; +import { readV3Bundle } from './bundle-v3'; import { sanitizePublishState } from '../../publish/publish-state-validator'; /** Key-value metadata key the Python writer stamps with the bundle's annotation format version. */ @@ -108,59 +111,90 @@ function readNumericColumnTypes(metadata: FileMetaData): Record { +function splitBundleParts(arrayBuffer: ArrayBuffer): (ArrayBuffer | null)[] { const uint8Array = new Uint8Array(arrayBuffer); const delimiterPositions = findBundleDelimiterPositions(uint8Array); - // 2 delimiters (core only), 3 (with settings), or 4 (settings + statistics). - if (delimiterPositions.length < 2 || delimiterPositions.length > 4) { + if (delimiterPositions.length < 2 || delimiterPositions.length > 5) { throw new Error( - `Expected 2 to 4 delimiters in parquetbundle, found ${delimiterPositions.length}`, + `Expected 2 to 5 delimiters in parquetbundle, found ${delimiterPositions.length}`, ); } - /** - * Copy out part `index` — part 0 starts at byte 0, every later part right after the - * preceding delimiter, and the final part runs to the end of the buffer. Order is - * fixed by the writer: annotations, projections metadata, projections, settings, - * statistics. Bounding each part by the *next* delimiter is what keeps a trailing part - * from being glued onto its predecessor's tail — without it, a 5-part bundle would hand - * the settings parser the statistics part too. Returns null for a zero-byte slot (an - * empty settings placeholder when a bundle carries statistics but no settings). - */ - const partAt = (index: number): ArrayBuffer | null => { - // Out of range must be null, not a slice: `delimiterPositions[index - 1]` is undefined, - // `undefined + 8` is NaN, and `subarray(NaN, len)` coerces NaN to 0 — returning the whole - // bundle as if it were one part. - if (index < 0 || index > delimiterPositions.length) return null; + // Part 0 starts at byte 0, every later part right after the preceding delimiter, and + // the final part runs to the end of the buffer. Bounding each part by the *next* + // delimiter is what keeps a trailing part from being glued onto its predecessor's + // tail — without it, a 5-part bundle would hand the settings parser the statistics + // part too. + const parts: (ArrayBuffer | null)[] = []; + for (let index = 0; index <= delimiterPositions.length; index++) { const view = uint8Array.subarray( index === 0 ? 0 : delimiterPositions[index - 1] + BUNDLE_DELIMITER_BYTES.length, index < delimiterPositions.length ? delimiterPositions[index] : uint8Array.length, ); - return view.byteLength > 0 ? view.slice().buffer : null; - }; + parts.push(view.byteLength > 0 ? view.slice().buffer : null); + } + return parts; +} + +/** + * Parse the annotations part's footer, or null when it is not readable parquet. + * + * Callers reuse the result both to read the format version and as the `metadata` + * option of the subsequent read — hyparquet re-derives metadata from the buffer when + * `metadata` is omitted, so passing it explicitly avoids parsing the same footer twice. + * A parse failure is swallowed here so the legacy reader keeps behaving exactly as it + * did: `formatVersion = 1`, and `parquetReadObjects` re-attempts the parse itself and + * surfaces the real error. + */ +function readPart1Metadata(part1: ArrayBuffer | null): FileMetaData | null { + if (!part1) return null; + try { + return parquetMetadata(part1); + } catch { + return null; + } +} + +/** + * Extract rows and optional settings from a parquetbundle (formats 1 and 2). + */ +export async function extractRowsFromParquetBundle( + arrayBuffer: ArrayBuffer, +): Promise { + const parts = splitBundleParts(arrayBuffer); + return extractRowsFromParts(parts, readPart1Metadata(parts[0])); +} - // The three required core parts. - let part1: ArrayBuffer | null = partAt(0); - let part2: ArrayBuffer | null = partAt(1); - let part3: ArrayBuffer | null = partAt(2); - const part4 = partAt(3); - const part5 = partAt(4); +/** + * The row-object reader for bundle formats 1 and 2, over parts already split out of + * the container. Split from {@link extractRowsFromParquetBundle} so the version sniff + * in {@link decodeParquetBundle} does not have to scan a 200 MB buffer twice. + */ +async function extractRowsFromParts( + parts: readonly (ArrayBuffer | null)[], + part1Metadata: FileMetaData | null, +): Promise { + let part1: ArrayBuffer | null = parts[0] ?? null; + let part2: ArrayBuffer | null = parts[1] ?? null; + let part3: ArrayBuffer | null = parts[2] ?? null; + const part4 = parts[3] ?? null; + const part5 = parts[4] ?? null; if (!part1 || !part2 || !part3) { throw new Error('Parquetbundle is missing one of its three required core parts'); @@ -171,22 +205,10 @@ export async function extractRowsFromParquetBundle( assertValidParquetMagic(part2); assertValidParquetMagic(part3); - // Parse part1's footer once (the annotations part), before it's decoded, and reuse - // the result both to read the format_version and as the `metadata` option below — - // hyparquet re-derives metadata from the buffer when `metadata` is omitted, so - // passing it explicitly avoids parsing the same footer twice. On parse failure, - // fall back to `formatVersion = 1` and let `parquetReadObjects` (without `metadata`) - // re-attempt the parse itself, surfacing the same error it would have before. - let part1Metadata: FileMetaData | null = null; - let formatVersion = 1; - let numericColumnTypes: Readonly> = {}; - try { - part1Metadata = parquetMetadata(part1); - formatVersion = readFormatVersion(part1Metadata); - numericColumnTypes = readNumericColumnTypes(part1Metadata); - } catch { - formatVersion = 1; - } + const formatVersion = part1Metadata ? readFormatVersion(part1Metadata) : 1; + const numericColumnTypes: Readonly> = part1Metadata + ? readNumericColumnTypes(part1Metadata) + : {}; // Decode sequentially and release each sliced buffer immediately after its decode completes. // hyparquet is CPU-bound on the single JS thread — Promise.all gives no real parallelism, only @@ -272,7 +294,7 @@ export async function extractRowsFromParquetBundle( * nothing here — a failed parse, an unmodelled column, a coerced type — can reach a file the * user saves. */ -async function extractStatistics( +export async function extractStatistics( statisticsBuffer: ArrayBuffer, ): Promise { try { @@ -309,7 +331,7 @@ async function extractStatistics( * Extract and parse settings from the 4th part of the bundle. * Returns null if parsing fails (graceful degradation). */ -async function extractSettings(settingsBuffer: ArrayBuffer): Promise { +export async function extractSettings(settingsBuffer: ArrayBuffer): Promise { try { // Validate parquet magic assertValidParquetMagic(settingsBuffer); @@ -345,6 +367,32 @@ async function extractSettings(settingsBuffer: ArrayBuffer): Promise { + const parts = splitBundleParts(arrayBuffer); + const part1Metadata = readPart1Metadata(parts[0]); + + if (part1Metadata && readFormatVersion(part1Metadata) >= 3) { + return readV3Bundle(parts, part1Metadata); + } + + const extraction = await extractRowsFromParts(parts, part1Metadata); + validateRowsBasic(extraction.projections); + return { + data: await convertParquetToVisualizationDataOptimized(extraction), + settings: extraction.settings, + }; +} + export function findColumn(columnNames: string[], candidates: string[]): string | null { for (const candidate of candidates) { const found = columnNames.find((col) => col.toLowerCase().includes(candidate.toLowerCase())); diff --git a/packages/core/src/components/data-loader/utils/conversion.ts b/packages/core/src/components/data-loader/utils/conversion.ts index 5cc4588c..4ebc351f 100644 --- a/packages/core/src/components/data-loader/utils/conversion.ts +++ b/packages/core/src/components/data-loader/utils/conversion.ts @@ -134,7 +134,7 @@ function* valuesForColumn(rows: Rows, column: string): Iterable { } } -function createNumericAnnotation( +export function createNumericAnnotation( numericType: 'int' | 'float', runtime?: Annotation['runtime'], ): Annotation { @@ -448,6 +448,30 @@ function appendSyntheticNACategory( } } +/** + * {@link appendSyntheticNACategory} for a dictionary-code column: missing slots are + * already `-1`, so the synthetic category is appended and every `-1` routed to it. + * + * Mutates the input arrays in place. Shared with the format v3 reader so both storage + * shapes gain the `__NA__` legend row under exactly one rule. + */ +export function appendSyntheticNACategoryToCodes( + uniqueValues: string[], + colors: string[], + shapes: string[], + codes: Int32Array, +): void { + if (!codes.some((code) => code < 0)) return; + + const naIndex = uniqueValues.length; + uniqueValues.push(NA_VALUE); + colors.push(NA_DEFAULT_COLOR); + shapes.push('circle'); + for (let p = 0; p < codes.length; p++) { + if (codes[p] < 0) codes[p] = naIndex; + } +} + /** * Parse an annotation value that may contain a pipe-separated score or evidence code suffix. * Format: `label|score`, `label|score1,score2,...`, or `label|EVIDENCE_CODE` @@ -633,7 +657,7 @@ function parseInfoJson(value: unknown): Record { * Builds a metadata map from projections metadata rows. * Parses info_json field and spreads its contents into metadata. */ -function buildProjectionsMetadataMap( +export function buildProjectionsMetadataMap( projectionsMetadata?: Rows, ): Map> { const metadataMap = new Map>(); @@ -775,9 +799,9 @@ export function convertParquetToVisualizationData( * and the parsed rows so the UI can render them. Raw `Rows` input (plain .parquet / legacy * reads) never carries either, so it passes straight through. */ -function carryStatistics( +export function carryStatistics( data: VisualizationData, - input: BundleExtractionResult | Rows, + input: Pick | Rows, ): VisualizationData { if (!Array.isArray(input) && input.statistics) { data.statistics = input.statistics; @@ -1700,19 +1724,7 @@ async function extractAnnotationsByProtein( if (annotationDataArray) { appendSyntheticNACategory(uniqueValues, colors, shapes, annotationDataArray); } else if (annotationDataTyped) { - // Int32Array missing slots are already -1; append NA category and remap -1. - const hasAnyMissing = annotationDataTyped.some((v) => v < 0); - if (hasAnyMissing) { - const naIndex = uniqueValues.length; - uniqueValues.push(NA_VALUE); - colors.push(NA_DEFAULT_COLOR); - shapes.push('circle'); - for (let p = 0; p < annotationDataTyped.length; p++) { - if (annotationDataTyped[p] < 0) { - annotationDataTyped[p] = naIndex; - } - } - } + appendSyntheticNACategoryToCodes(uniqueValues, colors, shapes, annotationDataTyped); } annotations[annotationCol] = createCategoricalAnnotation(uniqueValues, colors, shapes); From a4114a3d50459875c0820d0ed108a987285b6c30 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 01:33:24 +0200 Subject: [PATCH 17/31] perf(core): transfer the CSR and score payloads out of the decode worker --- .../data-loader/decode-transferables.test.ts | 82 +++++++++++++++++++ .../data-loader/decode-transferables.ts | 43 ++++++++++ .../components/data-loader/decode.worker.ts | 26 +----- 3 files changed, 128 insertions(+), 23 deletions(-) create mode 100644 packages/core/src/components/data-loader/decode-transferables.test.ts create mode 100644 packages/core/src/components/data-loader/decode-transferables.ts diff --git a/packages/core/src/components/data-loader/decode-transferables.test.ts b/packages/core/src/components/data-loader/decode-transferables.test.ts new file mode 100644 index 00000000..04d3cc47 --- /dev/null +++ b/packages/core/src/components/data-loader/decode-transferables.test.ts @@ -0,0 +1,82 @@ +import { describe, it, expect } from 'vitest'; +import type { VisualizationData } from '@protspace/utils'; +import { collectTransferables } from './decode-transferables'; + +/** + * Hand-built CSR dataset. `end` and `codes` deliberately share one ArrayBuffer, which + * is the case that makes deduplication load-bearing: `postMessage` throws + * `DataCloneError` on a transfer list that names the same buffer twice. + */ +function csrDataset(): { data: VisualizationData; shared: ArrayBuffer } { + const shared = new ArrayBuffer(6 * 4); + const end = new Int32Array(shared, 0, 3); // 3 proteins + const codes = new Int32Array(shared, 12, 3); + end.set([1, 2, 3]); + codes.set([0, 1, 0]); + + return { + shared, + data: { + protein_ids: ['P1', 'P2', 'P3'], + projections: [ + { name: 'pca2', data: new Float32Array(6), dimension: 2 }, + { name: 'umap3', data: new Float32Array(9), dimension: 3 }, + ], + annotations: { + go_bp: { kind: 'categorical', values: ['a', 'b'], colors: [], shapes: [] }, + organism: { kind: 'categorical', values: ['x'], colors: [], shapes: [] }, + }, + annotation_data: { + go_bp: { kind: 'csr', end, codes, length: 3 }, + organism: new Int32Array([0, 0, 0]), + }, + annotation_scores_csr: { + go_bp: { hitEnd: new Int32Array([1, 1, 2]), values: new Float32Array([0.5, 0.25]) }, + }, + annotation_evidence_csr: { + go_bp: { codes: new Int32Array([-1, 0, -1]), dict: ['IDA'] }, + }, + }, + }; +} + +describe('collectTransferables', () => { + it('names every bulk buffer exactly once, even when two views share one', () => { + const { data, shared } = csrDataset(); + const transfer = collectTransferables(data); + + expect(new Set(transfer).size).toBe(transfer.length); + expect(transfer).toContain(shared); + // 2 projections + the shared CSR buffer + organism codes + score hitEnd + score + // values + evidence codes. Without deduplication this would be 8: `end` and + // `codes` would each name `shared`. + expect(transfer).toHaveLength(7); + }); + + it('actually transfers: every source buffer is detached afterwards', () => { + const { data } = csrDataset(); + const transfer = collectTransferables(data); + const sources = [ + ...data.projections.map((projection) => projection.data), + data.annotation_data.organism as Int32Array, + data.annotation_scores_csr!.go_bp.values, + data.annotation_evidence_csr!.go_bp.codes, + ]; + + const clone = structuredClone(data, { transfer }); + + expect(clone.protein_ids).toEqual(['P1', 'P2', 'P3']); + expect(sources.every((array) => array.byteLength === 0)).toBe(true); + }); + + it('leaves a v1/v2 dataset with only its projection and Int32Array buffers', () => { + const data: VisualizationData = { + protein_ids: ['P1'], + projections: [{ name: 'pca2', data: new Float32Array(2), dimension: 2 }], + annotations: { organism: { kind: 'categorical', values: ['x'], colors: [], shapes: [] } }, + annotation_data: { organism: new Int32Array([0]), multi: [[0]] }, + }; + + expect(collectTransferables(data)).toHaveLength(2); + }); +}); diff --git a/packages/core/src/components/data-loader/decode-transferables.ts b/packages/core/src/components/data-loader/decode-transferables.ts new file mode 100644 index 00000000..cf1be846 --- /dev/null +++ b/packages/core/src/components/data-loader/decode-transferables.ts @@ -0,0 +1,43 @@ +import { isCsrAnnotationData, type VisualizationData } from '@protspace/utils'; + +/** + * ArrayBuffers of a decoded dataset that can move to the main thread zero-copy. + * + * Everything a v3 read produces in bulk is a typed array, and structured-cloning those + * is what the format was designed to avoid: at 573K proteins the clone of the result + * alone cost 3.7 s. What is left behind — the id strings, the projection metadata, the + * numeric `(number | null)[]` columns — is cloned as before. + * + * Deduplicated through a `Set`, because `postMessage` throws on a duplicate entry in + * the transfer list, and two views can legitimately share one buffer. + * + * Lives outside `decode.worker.ts` so node tests can import it: that module is only + * loadable through Vite's `?worker` transform. + */ +export function collectTransferables(data: VisualizationData): Transferable[] { + const buffers = new Set(); + const push = (buffer: ArrayBufferLike | undefined) => { + if (buffer instanceof ArrayBuffer) buffers.add(buffer); + }; + + for (const projection of data.projections) { + if (projection.data instanceof Float32Array) push(projection.data.buffer); + } + for (const value of Object.values(data.annotation_data)) { + if (value instanceof Int32Array) { + push(value.buffer); + } else if (isCsrAnnotationData(value)) { + push(value.end.buffer); + push(value.codes.buffer); + } + } + for (const scores of Object.values(data.annotation_scores_csr ?? {})) { + push(scores.hitEnd.buffer); + push(scores.values.buffer); + } + for (const evidence of Object.values(data.annotation_evidence_csr ?? {})) { + push(evidence.codes.buffer); + } + + return [...buffers]; +} diff --git a/packages/core/src/components/data-loader/decode.worker.ts b/packages/core/src/components/data-loader/decode.worker.ts index b5504a47..db45f6e3 100644 --- a/packages/core/src/components/data-loader/decode.worker.ts +++ b/packages/core/src/components/data-loader/decode.worker.ts @@ -1,7 +1,5 @@ -import { extractRowsFromParquetBundle } from './utils/bundle'; -import { convertParquetToVisualizationDataOptimized } from './utils/conversion'; -import { validateRowsBasic } from './utils/validation'; -import type { VisualizationData, BundleSettings } from '@protspace/utils'; +import { decodeParquetBundle } from './utils/bundle'; +import { collectTransferables } from './decode-transferables'; interface DecodeRequest { type: 'decode-bundle'; @@ -13,28 +11,10 @@ const ctx = self as unknown as { postMessage(message: unknown, transfer: Transferable[]): void; }; -/** - * Collect transferable ArrayBuffers from the result (Float32 coords + Int32 annotation columns) - * so they move zero-copy. Everything else (strings, metadata, number[][]) is structured-cloned. - */ -function collectTransferables(data: VisualizationData): Transferable[] { - const transfer: Transferable[] = []; - for (const projection of data.projections) { - if (projection.data instanceof Float32Array) transfer.push(projection.data.buffer); - } - for (const value of Object.values(data.annotation_data)) { - if (value instanceof Int32Array) transfer.push(value.buffer); - } - return transfer; -} - ctx.onmessage = async (event: MessageEvent) => { const { arrayBuffer } = event.data; try { - const extraction = await extractRowsFromParquetBundle(arrayBuffer); - validateRowsBasic(extraction.projections); - const data = await convertParquetToVisualizationDataOptimized(extraction); - const settings: BundleSettings | null = extraction.settings; + const { data, settings } = await decodeParquetBundle(arrayBuffer); ctx.postMessage( { type: 'decode-result', ok: true, data, settings }, collectTransferables(data), From ab97b94c8494e8d73c30086805bd7b79a4cb39a2 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 02:06:13 +0200 Subject: [PATCH 18/31] fix(core): validate the v3 manifest against part 1's physical column types Bound num_rows, range-check evidence codes, reject id and payload-name collisions. --- .../data-loader/decode-transferables.test.ts | 42 +++- .../data-loader/utils/bundle-v3.test.ts | 180 ++++++++++++++++-- .../components/data-loader/utils/bundle-v3.ts | 104 ++++++++-- .../components/data-loader/utils/bundle.ts | 9 + 4 files changed, 294 insertions(+), 41 deletions(-) diff --git a/packages/core/src/components/data-loader/decode-transferables.test.ts b/packages/core/src/components/data-loader/decode-transferables.test.ts index 04d3cc47..c4ba2c9c 100644 --- a/packages/core/src/components/data-loader/decode-transferables.test.ts +++ b/packages/core/src/components/data-loader/decode-transferables.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect } from 'vitest'; -import type { VisualizationData } from '@protspace/utils'; +import { isCsrAnnotationData, type VisualizationData } from '@protspace/utils'; import { collectTransferables } from './decode-transferables'; /** @@ -40,6 +40,23 @@ function csrDataset(): { data: VisualizationData; shared: ArrayBuffer } { }; } +/** Every typed array `collectTransferables` names a buffer for, in a stable order. */ +const bulkViews = (data: VisualizationData): (Int32Array | Float32Array)[] => [ + ...data.projections.map((projection) => projection.data as Float32Array), + ...Object.values(data.annotation_data).flatMap((value) => + value instanceof Int32Array + ? [value] + : isCsrAnnotationData(value) + ? [value.end, value.codes] + : [], + ), + ...Object.values(data.annotation_scores_csr ?? {}).flatMap((scores) => [ + scores.hitEnd, + scores.values, + ]), + ...Object.values(data.annotation_evidence_csr ?? {}).map((evidence) => evidence.codes), +]; + describe('collectTransferables', () => { it('names every bulk buffer exactly once, even when two views share one', () => { const { data, shared } = csrDataset(); @@ -53,20 +70,29 @@ describe('collectTransferables', () => { expect(transfer).toHaveLength(7); }); - it('actually transfers: every source buffer is detached afterwards', () => { + it('actually transfers: the clone holds the bytes and every source is detached', () => { const { data } = csrDataset(); const transfer = collectTransferables(data); - const sources = [ - ...data.projections.map((projection) => projection.data), - data.annotation_data.organism as Int32Array, - data.annotation_scores_csr!.go_bp.values, - data.annotation_evidence_csr!.go_bp.codes, - ]; + const sources = bulkViews(data); + const before = sources.map((view) => Array.from(view)); const clone = structuredClone(data, { transfer }); expect(clone.protein_ids).toEqual(['P1', 'P2', 'P3']); expect(sources.every((array) => array.byteLength === 0)).toBe(true); + // A detached sender proves only that something moved. What has to survive is the + // content, including the two views that share one buffer at different offsets. + expect(bulkViews(clone).map((view) => Array.from(view))).toEqual(before); + expect(before).toEqual([ + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 2, 3], + [0, 1, 0], + [0, 0, 0], + [1, 1, 2], + [0.5, 0.25], + [-1, 0, -1], + ]); }); it('leaves a v1/v2 dataset with only its projection and Int32Array buffers', () => { diff --git a/packages/core/src/components/data-loader/utils/bundle-v3.test.ts b/packages/core/src/components/data-loader/utils/bundle-v3.test.ts index 41128236..7953af05 100644 --- a/packages/core/src/components/data-loader/utils/bundle-v3.test.ts +++ b/packages/core/src/components/data-loader/utils/bundle-v3.test.ts @@ -1,6 +1,7 @@ import { describe, it, expect, vi, afterEach } from 'vitest'; import { readFileSync } from 'node:fs'; import { parquetWriteBuffer } from 'hyparquet-writer'; +import { parquetMetadata } from 'hyparquet'; import { BUNDLE_DELIMITER_BYTES, concatenateBuffers, @@ -12,7 +13,8 @@ import { type CsrAnnotationData, type VisualizationData, } from '@protspace/utils'; -import { decodeParquetBundle } from './bundle'; +import { decodeParquetBundle, extractRowsFromParquetBundle } from './bundle'; +import { readV3Bundle } from './bundle-v3'; import { collectTransferables } from '../decode-transferables'; /** @@ -143,6 +145,26 @@ const v3Bundle = (overrides: Record = {}) => ), ); +/** + * Every typed array `collectTransferables` names a buffer for, in a stable order, so a + * dataset and its structured clone can be compared element by element. + */ +const bulkViews = (data: VisualizationData): (Int32Array | Float32Array)[] => [ + ...data.projections.map((projection) => projection.data as Float32Array), + ...Object.values(data.annotation_data).flatMap((value) => + value instanceof Int32Array + ? [value] + : isCsrAnnotationData(value) + ? [value.end, value.codes] + : [], + ), + ...Object.values(data.annotation_scores_csr ?? {}).flatMap((scores) => [ + scores.hitEnd, + scores.values, + ]), + ...Object.values(data.annotation_evidence_csr ?? {}).map((evidence) => evidence.codes), +]; + const labelsOf = (data: VisualizationData, key: string, protein: number) => getProteinAnnotationIndices(data.annotation_data[key], protein).map( (index) => data.annotations[key].values[index], @@ -322,6 +344,32 @@ describe('parquetbundle format v3', () => { { ...MANIFEST, projections: [{ name: 'nope', dimension: 2 }] }, /part 3 has no nope__x/, ], + [ + 'a code column the manifest calls numeric', + { ...MANIFEST, columns: { ...MANIFEST.columns, organism: { kind: 'numeric' } } }, + /is kind "numeric", but part 1 stores "organism" as INT32, not DOUBLE/, + ], + [ + 'a numeric column the manifest calls categorical', + { ...MANIFEST, columns: { ...MANIFEST.columns, score: { kind: 'categorical' } } }, + /is kind "categorical", but part 1 stores "score" as DOUBLE, not INT32/, + ], + [ + 'a hit-count column the manifest calls numeric', + { + ...MANIFEST, + columns: { ...MANIFEST.columns, go_bp__count: { kind: 'numeric' } }, + }, + /is kind "numeric", but part 1 stores "go_bp__count" as INT32, not DOUBLE/, + ], + [ + 'an annotation column that collides with the id column', + { + ...MANIFEST, + columns: { ...MANIFEST.columns, protein_id: { kind: 'numeric' } }, + }, + /declares idColumn "protein_id" as an annotation column too/, + ], [ 'an unknown numericType', { @@ -398,6 +446,33 @@ describe('parquetbundle format v3', () => { ); }); + it('an evidence code outside the evidence dictionary', async () => { + const payloads = { ...PAYLOADS, 'evidence:go_bp': i32(-1, 0, 1, -1, -1, -1, 5, -1, -1) }; + await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( + /hit 6 has evidence code 5, outside the 2 evidence labels/, + ); + }); + + it('an evidence code below the -1 "no evidence" sentinel', async () => { + const payloads = { ...PAYLOADS, 'evidence:go_bp': i32(-2, 0, 1, -1, -1, -1, 0, -1, -1) }; + await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( + /hit 0 has evidence code -2, outside the 2 evidence labels/, + ); + }); + + it('two payloads sharing one name', async () => { + // `payloadPart` takes a Record, which cannot hold a duplicate key, so the rows + // are written directly. + const names = [...Object.keys(PAYLOADS), 'csr:go_bp']; + const duplicated = part([ + { name: 'name', data: names }, + { name: 'data', data: [...Object.values(PAYLOADS), i32(0, 0, 0, 0, 0, 0, 0, 0, 0)] }, + ]); + await expect(decodeParquetBundle(v3Bundle({ 5: duplicated }))).rejects.toThrow( + /payloads part declares "csr:go_bp" twice/, + ); + }); + it('a missing payloads part', async () => { await expect(decodeParquetBundle(v3Bundle({ 5: EMPTY }))).rejects.toThrow( /carries no payloads part/, @@ -405,6 +480,55 @@ describe('parquetbundle format v3', () => { }); }); + // `parquetWriteBuffer` always stamps a truthful `num_rows`, so the lying footer is + // built by handing `readV3Bundle` a doctored `FileMetaData` — the same object + // `decodeParquetBundle` reads out of part 1. + it.each([ + ['above the row cap', 2_000_001n], + ['negative', -1n], + ['past the safe-integer range', 9_007_199_254_740_993n], + ['absent', undefined], + ])('rejects a footer whose row count is %s before allocating on it', async (_label, rows) => { + const part1 = annotationsPart(); + const parts = [ + part1, + PROJECTIONS_METADATA, + PROJECTIONS, + EMPTY, + EMPTY, + payloadPart(PAYLOADS), + ].map((buffer) => (buffer.byteLength > 0 ? (buffer.slice().buffer as ArrayBuffer) : null)); + const metadata = parquetMetadata(parts[0]!); + + await expect(readV3Bundle(parts, { ...metadata, num_rows: rows as bigint })).rejects.toThrow( + /rows, outside 0\.\.2000000/, + ); + }); + + it('reports a projection column that was not written REQUIRED and PLAIN', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const nullable = new Uint8Array( + parquetWriteBuffer({ + columnData: [ + { name: 'pca2__x', data: [1, 2, 3, 4, 5, 6, 7, null], type: 'FLOAT', nullable: true }, + { name: 'pca2__y', data: new Float32Array([1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5]) }, + { name: 'umap3__x', data: new Float32Array([10, 20, 30, 40, 50, 60, 0, 0]) }, + { name: 'umap3__y', data: new Float32Array([11, 21, 31, 41, 51, 61, 0, 0]) }, + { name: 'umap3__z', data: new Float32Array([0.25, 0.5, 0.75, 1, 1.25, 1.5, 0, 0]) }, + ] as never, + statistics: false, + }), + ); + + const { data } = await decodeParquetBundle(v3Bundle({ 2: nullable })); + + // The null coerces to 0, which is indistinguishable from an absent protein's + // origin fallback — the warning is the only signal that it happened. + expect(Array.from(data.projections[0].data.slice(14))).toEqual([0, 8.5]); + expect(warn).toHaveBeenCalledTimes(1); + expect(warn.mock.calls[0][0]).toMatch(/column "pca2__x" did not decode to a typed array/); + }); + it('decodes non-ASCII labels by byte range, not character offset', async () => { // 'Mü' is three bytes but two characters, so slicing the decoded blob by byte // offsets would shear every later label. @@ -417,6 +541,26 @@ describe('parquetbundle format v3', () => { expect(data.annotations.organism.values).toEqual(['Human', 'Mü', 'Yeast', 'Fly', NA_VALUE]); }); + it('refuses to read a v3 bundle through the legacy row-object extractor', async () => { + // Widening the delimiter gate to 5 made this reachable: without the version guard + // it gets as far as part 3 and complains about missing projection columns. + await expect(extractRowsFromParquetBundle(v3Bundle())).rejects.toThrow( + /declares annotation format v3, which only decodeParquetBundle can read/, + ); + }); + + it('keeps a byte-order mark that belongs to a label', async () => { + // U+FEFF is three bytes, and a decoder that treats it as an encoding marker rather + // than a character silently renames the category. + const payloads = { + ...PAYLOADS, + 'dict:organism': utf8('\uFEFFHumanMouseYeastFly'), + 'dict:organism:len': i32(8, 5, 5, 3), + }; + const { data } = await decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) })); + expect(data.annotations.organism.values[0]).toBe('\uFEFFHuman'); + }); + it('still reads a bundle whose columns were written nullable, and says so once', async () => { const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}); const nullable = new Uint8Array( @@ -466,18 +610,30 @@ describe('parquetbundle format v3', () => { const transfer = collectTransferables(data); expect(new Set(transfer).size).toBe(transfer.length); - // 2 projections + organism codes + 2 CSR (end + codes) + scores (hitEnd + values) - // + evidence codes. + // 2 projections + organism codes + 2 x 2 CSR (end + codes) + scores (hitEnd + + // values) + evidence codes. expect(transfer).toHaveLength(10); - const sources = [ - ...data.projections.map((projection) => projection.data), - data.annotation_data.organism as Int32Array, - (data.annotation_data.go_bp as CsrAnnotationData).codes, - data.annotation_scores_csr!.go_bp.values, - data.annotation_evidence_csr!.go_bp.codes, - ]; - structuredClone(data, { transfer }); - expect(sources.every((array) => array.byteLength === 0)).toBe(true); + const sources = bulkViews(data); + const before = sources.map((view) => Array.from(view)); + + // Each transferred buffer must be owned outright by exactly one view. A view into a + // slice of someone else's buffer (a hyparquet page, say) would still transfer, but it + // would carry — and detach — bytes that are not ours. + for (const buffer of transfer as ArrayBuffer[]) { + const owners = sources.filter((view) => view.buffer === buffer); + expect(owners).toHaveLength(1); + expect(owners[0].byteOffset).toBe(0); + expect(owners[0].byteLength).toBe(buffer.byteLength); + } + + const clone = structuredClone(data, { transfer }); + + expect(sources.every((view) => view.byteLength === 0)).toBe(true); + // Reading the clone is the point: asserting only that the sender detached would + // pass just as happily on a clone holding the wrong bytes. + expect(bulkViews(clone).map((view) => Array.from(view))).toEqual(before); + expect(clone.protein_ids).toEqual(PROTEIN_IDS); + expect(clone.annotations.go_bp.values).toEqual(data.annotations.go_bp.values); }); }); diff --git a/packages/core/src/components/data-loader/utils/bundle-v3.ts b/packages/core/src/components/data-loader/utils/bundle-v3.ts index 074e3d9b..e423e7cb 100644 --- a/packages/core/src/components/data-loader/utils/bundle-v3.ts +++ b/packages/core/src/components/data-loader/utils/bundle-v3.ts @@ -35,7 +35,7 @@ import { type Projection, type VisualizationData, } from '@protspace/utils'; -import { assertValidParquetMagic } from './validation'; +import { assertValidParquetMagic, DEFAULT_VALIDATION_LIMITS } from './validation'; import { extractSettings, extractStatistics } from './bundle'; import { appendSyntheticNACategoryToCodes, @@ -55,7 +55,9 @@ const EVIDENCE_DICT_NAME = '__evidence'; const AXES = ['x', 'y', 'z'] as const; -const DECODER = new TextDecoder(); +// ignoreBOM keeps a leading U+FEFF as a character: it is part of a label, not an +// encoding marker, and stripping it would silently rename the category. +const DECODER = new TextDecoder('utf-8', { ignoreBOM: true }); type V3ColumnKind = 'categorical' | 'multi' | 'numeric'; @@ -80,15 +82,22 @@ function physicalColumn(name: string, kind: V3ColumnKind): string { return kind === 'multi' ? `${name}__count` : name; } -/** Leaf (data) column names of a parquet schema; the root element carries no type. */ -function leafColumnNames(metadata: FileMetaData): Set { - const names = new Set(); +/** Leaf (data) columns of a parquet schema as `name -> physical type`; the root carries no type. */ +function leafColumnTypes(metadata: FileMetaData): Map { + const types = new Map(); for (const field of metadata.schema) { - if (field.name && field.type) names.add(field.name); + if (field.name && field.type) types.set(field.name, field.type); } - return names; + return types; } +/** Physical parquet type the encoder writes for each kind, and the reader assumes. */ +const PHYSICAL_TYPE: Record = { + numeric: 'DOUBLE', + categorical: 'INT32', + multi: 'INT32', +}; + function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value); } @@ -116,7 +125,7 @@ function readManifest(metadata: FileMetaData): V3Manifest { } if (!isRecord(parsed)) throw new Error('v3 manifest is not a JSON object'); - const schemaColumns = leafColumnNames(metadata); + const schemaColumns = leafColumnTypes(metadata); const { idColumn, columns, projections } = parsed; if (typeof idColumn !== 'string' || !schemaColumns.has(idColumn)) { @@ -138,10 +147,24 @@ function readManifest(metadata: FileMetaData): V3Manifest { `v3 manifest column "${name}" has unknown numericType "${String(numericType)}"`, ); } + if (name === idColumn) { + throw new Error(`v3 manifest declares idColumn "${name}" as an annotation column too`); + } const physical = physicalColumn(name, kind); - if (!schemaColumns.has(physical)) { + const physicalType = schemaColumns.get(physical); + if (physicalType === undefined) { throw new Error(`v3 manifest declares column "${name}" but part 1 has no "${physical}"`); } + // The kind is the ONLY thing that says how the stored numbers are read, so it is + // checked against what they physically are: the encoder writes every numeric as + // float64 and every dictionary code / hit count as int32. Without this, a manifest + // calling a code column numeric turns dictionary codes into a colour gradient. + if (physicalType !== PHYSICAL_TYPE[kind]) { + throw new Error( + `v3 manifest column "${name}" is kind "${kind}", but part 1 stores "${physical}" ` + + `as ${physicalType}, not ${PHYSICAL_TYPE[kind]}`, + ); + } validated[name] = { kind, ...(numericType != null ? { numericType } : {}), @@ -213,6 +236,25 @@ function writeChunk( } } +/** + * One-shot reporter for a column that did not arrive as a typed array. + * + * Per read rather than per column: a producer that got this wrong got it wrong for the + * whole part, and one line in the console is the point. + */ +function plainArrayReporter(): (columnName: string) => void { + let warned = false; + return (columnName: string) => { + if (warned) return; + warned = true; + console.warn( + `v3 bundle column "${columnName}" did not decode to a typed array — it was probably ` + + 'written nullable or dictionary-encoded. The bundle still loads, about 4x slower; ' + + 'fix the writer (every v3 column must be REQUIRED and PLAIN).', + ); + }; +} + /** Preallocate one array per declared column and fill it chunk by chunk. */ async function readAnnotationColumns( part: ArrayBuffer, @@ -229,16 +271,7 @@ async function readAnnotationColumns( ); } - let warned = false; - const onPlainArray = (columnName: string) => { - if (warned) return; - warned = true; - console.warn( - `v3 bundle column "${columnName}" did not decode to a typed array — it was probably ` + - 'written nullable or dictionary-encoded. The bundle still loads, about 4x slower; ' + - 'fix the writer (every v3 column must be REQUIRED and PLAIN).', - ); - }; + const onPlainArray = plainArrayReporter(); await parquetRead({ file: part, @@ -269,7 +302,7 @@ async function readProjections( ): Promise { assertValidParquetMagic(part); const metadata = parquetMetadata(part); - const schemaColumns = leafColumnNames(metadata); + const schemaColumns = leafColumnTypes(metadata); const axisTargets = new Map(); const projections: Projection[] = []; @@ -294,6 +327,7 @@ async function readProjections( } if (axisTargets.size > 0) { + const onPlainArray = plainArrayReporter(); await parquetRead({ file: part, metadata, @@ -301,6 +335,9 @@ async function readProjections( onChunk: ({ columnName, columnData, rowStart }) => { const target = axisTargets.get(columnName); if (!target) return; + // A nullable axis column coerces its nulls to 0 below, which is exactly the + // origin an absent protein legitimately sits at — so it has to be reported. + if (Array.isArray(columnData)) onPlainArray(columnName); const { data, dimension, axis } = target; for (let i = 0; i < columnData.length; i++) { data[(rowStart + i) * dimension + axis] = columnData[i] as number; @@ -321,6 +358,9 @@ async function readPayloads(part: ArrayBuffer): Promise> const payloads = new Map(); for (const row of rows) { const name = typeof row.name === 'string' ? row.name : DECODER.decode(row.name as Uint8Array); + // Last-win would silently pick one of two disagreeing payloads; the encoder already + // rejects the collision, so reaching here means a producer bug. + if (payloads.has(name)) throw new Error(`v3 payloads part declares "${name}" twice`); payloads.set(name, row.data as Uint8Array); } return payloads; @@ -453,7 +493,17 @@ function readCsrColumn( `v3 column "${name}" has ${evidenceCodes.length} evidence codes for ${codes.length} hits`, ); } - evidence = { codes: evidenceCodes, dict: evidenceDict() }; + const dict = evidenceDict(); + for (let hit = 0; hit < evidenceCodes.length; hit++) { + // -1 is "no evidence", which the reader itself writes for an inserted NA hit. + if (evidenceCodes[hit] < -1 || evidenceCodes[hit] >= dict.length) { + throw new Error( + `v3 column "${name}" hit ${hit} has evidence code ${evidenceCodes[hit]}, ` + + `outside the ${dict.length} evidence labels`, + ); + } + } + evidence = { codes: evidenceCodes, dict }; } return { end, codes, scores, evidence }; @@ -538,7 +588,19 @@ export async function readV3Bundle( } const manifest = readManifest(metadata); + // Everything below preallocates on this footer field before a single row is read, so + // it is bounded here. The v3 path never reaches `validateRowsBasic`, which is what + // caps the legacy path. const numRows = Number(metadata.num_rows); + if ( + !Number.isSafeInteger(numRows) || + numRows < 0 || + numRows > DEFAULT_VALIDATION_LIMITS.maxRows + ) { + throw new Error( + `v3 bundle declares ${String(metadata.num_rows)} rows, outside 0..${DEFAULT_VALIDATION_LIMITS.maxRows}`, + ); + } const columns = await readAnnotationColumns(part1, metadata, manifest, numRows); const protein_ids = columns.get(manifest.idColumn) as string[]; diff --git a/packages/core/src/components/data-loader/utils/bundle.ts b/packages/core/src/components/data-loader/utils/bundle.ts index 57ae9e67..be7dd68d 100644 --- a/packages/core/src/components/data-loader/utils/bundle.ts +++ b/packages/core/src/components/data-loader/utils/bundle.ts @@ -206,6 +206,15 @@ async function extractRowsFromParts( assertValidParquetMagic(part3); const formatVersion = part1Metadata ? readFormatVersion(part1Metadata) : 1; + // v3 stores its annotations as dictionary codes plus payloads, which this row-object + // reader cannot make sense of: it would get as far as part 3 and complain about + // missing 'projection_name'/'x'/'y' columns. Say what is actually wrong instead. + if (formatVersion >= 3) { + throw new Error( + `Parquetbundle declares annotation format v${formatVersion}, which only ` + + 'decodeParquetBundle can read; extractRowsFromParquetBundle handles v1 and v2.', + ); + } const numericColumnTypes: Readonly> = part1Metadata ? readNumericColumnTypes(part1Metadata) : {}; From 3464c3f339415bc5aca427516edefb0c927cddd1 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 02:06:30 +0200 Subject: [PATCH 19/31] perf(core): return the empty legend before its O(N) bincount Name the one in-place dataset writer in the hash memo's soundness note. --- packages/core/src/components/legend/legend.ts | 18 ++++++++++-------- packages/utils/src/storage/data-hash.ts | 2 ++ 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/packages/core/src/components/legend/legend.ts b/packages/core/src/components/legend/legend.ts index 22bcc9c7..e518bc38 100644 --- a/packages/core/src/components/legend/legend.ts +++ b/packages/core/src/components/legend/legend.ts @@ -1818,15 +1818,17 @@ export class ProtspaceLegend extends LitElement { // Aligned with PersistenceController's isNumericAnnotation callback so the // processor and the persistence layer agree on numeric-ness in transient states. const isNumericAnnotation = this._isCurrentAnnotationNumeric(); - const knownValues = - isNumericAnnotation && this.annotationData?.values?.length - ? this.annotationData.values.map((value) => toInternalValue(value)) - : []; + // Ahead of the bincount below, which is O(proteins) and has nothing to count for + // an annotation with no values at all. + if (!this.annotationData?.values?.length) { + this._legendItems = []; + return; + } + const knownValues = isNumericAnnotation + ? this.annotationData.values.map((value) => toInternalValue(value)) + : []; const frequencies = this._computeAnnotationCounts(knownValues); - if ( - !this.annotationData?.values?.length || - (!isNumericAnnotation && !(frequencies?.size ?? this.annotationValues?.length)) - ) { + if (!isNumericAnnotation && !(frequencies?.size ?? this.annotationValues?.length)) { this._legendItems = []; return; } diff --git a/packages/utils/src/storage/data-hash.ts b/packages/utils/src/storage/data-hash.ts index f87c3373..f3badf6c 100644 --- a/packages/utils/src/storage/data-hash.ts +++ b/packages/utils/src/storage/data-hash.ts @@ -244,6 +244,8 @@ function buildDatasetFingerprint(data: DatasetHashInput): string { * sound because no producer mutates a live dataset in place: every transform * (`materializeVisualizationData`, `cloneWithPredictions`, the conversion * pipeline) hands back fresh containers, which miss the memo and recompute. + * The one in-place writer is `restoreDeclaredNumericAnnotations` (conversion.ts), + * which runs inside that pipeline before any hash is taken — keep it there. */ interface DatasetHashMemo { annotations: DatasetHashInput['annotations']; From 2ac5bf724896cc7dc094f97e45937fe794c79fef Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 02:13:42 +0200 Subject: [PATCH 20/31] fix(bundle): make the legacy annotation migration idempotent and stamped An unstamped v2 output reads back as v1, so the next caller double-escaped it. --- .../protspace/data/annotations/encoding.py | 14 +++++++-- apps/protspace/tests/test_bundle_v3_encode.py | 30 ++++++++++++++++++- 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/apps/protspace/src/protspace/data/annotations/encoding.py b/apps/protspace/src/protspace/data/annotations/encoding.py index 889acedc..c5760696 100644 --- a/apps/protspace/src/protspace/data/annotations/encoding.py +++ b/apps/protspace/src/protspace/data/annotations/encoding.py @@ -135,7 +135,17 @@ def encode_legacy_cell(value: str) -> str: def migrate_legacy_annotation_table(table: pa.Table) -> pa.Table: - """Re-emit every v1 string annotation using unambiguous v2 field encoding.""" + """Re-emit every v1 string annotation using unambiguous v2 field encoding. + + Idempotent, and it has to be: :func:`read_format_version` reads an unstamped + table as v1, so an unstamped output is exactly the table that looks like it + still needs migrating, and a second pass double-escapes every reserved + character (``%3B`` to ``%253B``, unrecoverable because :func:`decode_field` + is not its own inverse). An already-migrated table is returned untouched and + the result is stamped, so neither the caller nor the next one can repeat it. + """ + if read_format_version(table) >= BUNDLE_FORMAT_VERSION: + return table columns = [] for name, column in zip(table.column_names, table.columns, strict=True): if name in {"identifier", "protein_id"} or not ( @@ -153,7 +163,7 @@ def migrate_legacy_annotation_table(table: pa.Table) -> pa.Table: for value in column.to_pylist() ] columns.append(pa.array(migrated, type=column.type)) - return pa.Table.from_arrays(columns, names=table.column_names) + return stamp_format_version(pa.Table.from_arrays(columns, names=table.column_names)) def to_display_value(raw, *, decode: bool = True): diff --git a/apps/protspace/tests/test_bundle_v3_encode.py b/apps/protspace/tests/test_bundle_v3_encode.py index 981e07fb..af0b9f71 100644 --- a/apps/protspace/tests/test_bundle_v3_encode.py +++ b/apps/protspace/tests/test_bundle_v3_encode.py @@ -18,7 +18,12 @@ import pyarrow.parquet as pq import pytest -from protspace.data.annotations.encoding import FORMAT_VERSION_KEY, stamp_format_version +from protspace.data.annotations.encoding import ( + FORMAT_VERSION_KEY, + migrate_legacy_annotation_table, + read_format_version, + stamp_format_version, +) from protspace.data.io.bundle_v3 import MANIFEST_KEY, encode_v3 @@ -471,6 +476,29 @@ def test_v1_annotations_are_migrated_before_splitting(): assert read(parts[0]).column("col__count").to_pylist() == [1] +def test_migrating_a_v1_table_twice_is_a_no_op(): + """The migration output must not read back as v1, or it gets re-escaped. + + ``read_format_version`` defaults an unstamped table to 1, so an unstamped + migration output is exactly the one table that looks like it still needs + migrating. ``decode_field`` is not its own inverse, so the second pass is + unrecoverable. + """ + v1 = pa.table({"protein_id": ["p0"], "col": ["nitrite reductase (a; b)"]}) + once = migrate_legacy_annotation_table(v1) + assert read_format_version(once) == 2 + assert migrate_legacy_annotation_table(once).equals(once) + + +def test_the_encoder_does_not_re_migrate_an_already_migrated_table(): + """``35K_ec_brenda`` row 28982, the cell that reproduced the double escape.""" + name = "nitrite reductase (cytochrome; ammonia-forming)" + migrated = migrate_legacy_annotation_table( + pa.table({"protein_id": ["p0"], "col": [name]}) + ) + assert labels_of(payloads_of(encode(migrated)[3]), "col") == [name] + + def test_null_cells_are_missing(): table = stamp_format_version( pa.table({"protein_id": ["p0", "p1", "p2"], "col": ["A", None, "A"]}) From 7119f8b76e7c0aa9e19c51e033f2b334937ebda7 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 02:14:59 +0200 Subject: [PATCH 21/31] fix(bundle): store v3 scores as float64 and keep missing-token spellings float32 flushed E-values to 0 or inf; 'none'/'NA' cells are labels, not gaps. --- .../src/protspace/data/io/bundle_v3.py | 56 +++++++++++++------ apps/protspace/tests/test_bundle_v3_decode.py | 25 +++++++-- apps/protspace/tests/test_bundle_v3_encode.py | 52 +++++++++++++++-- 3 files changed, 106 insertions(+), 27 deletions(-) diff --git a/apps/protspace/src/protspace/data/io/bundle_v3.py b/apps/protspace/src/protspace/data/io/bundle_v3.py index 136ad49e..d5f3fe6b 100644 --- a/apps/protspace/src/protspace/data/io/bundle_v3.py +++ b/apps/protspace/src/protspace/data/io/bundle_v3.py @@ -49,9 +49,16 @@ CONTAINER_VERSION = 3 MANIFEST_KEY = b"protspace_v3_manifest" -#: Cell/hit spellings that mean "missing". Mirrors ``MISSING_VALUE_TOKENS`` -#: in ``packages/utils/src/visualization/missing-values.ts``; compared against -#: the lower-cased, whitespace-trimmed token. +#: Cell spellings that block numeric inference, mirroring +#: ``MISSING_VALUE_TOKENS`` in +#: ``packages/utils/src/visualization/missing-values.ts``; compared against the +#: lower-cased, whitespace-trimmed cell. They are *only* consulted there: a +#: column of ``NA`` stays categorical instead of becoming all-NaN numeric, but a +#: cell literally spelled ``none`` keeps that label, because the file has to +#: preserve the token it was given (``protspace style`` and the Dash legend key +#: on it, and ``phosphatase.predicted_transmembrane`` is 1383 of 1587 rows of +#: literal ``none``). The browser re-applies ``normalizeMissingValue`` at read +#: time, so folding these into NA stays *its* decision, on both v2 and v3. MISSING_TOKENS = frozenset({"na", "n/a", "nan", "null", "none", "__na__"}) #: ``EVIDENCE_CODE_RE`` from ``conversion.ts``: the part after a hit's last @@ -145,13 +152,16 @@ def _as_string(column: pa.ChunkedArray | pa.Array) -> pa.Array: return pc.cast(arr, pa.string()) +def _blank_mask(trimmed: pa.Array) -> np.ndarray: + """Genuinely absent: null or the empty string (whitespace already trimmed).""" + mask = pc.or_(pc.is_null(trimmed), pc.equal(trimmed, pa.scalar(""))) + return np.asarray(pc.fill_null(mask, True)) + + def _missing_mask(trimmed: pa.Array) -> np.ndarray: """``normalizeMissingValue``: null, blank, or a MISSING_TOKENS spelling.""" - is_null = pc.is_null(trimmed) - blank = pc.equal(trimmed, pa.scalar("")) token = pc.is_in(pc.utf8_lower(trimmed), value_set=pa.array(sorted(MISSING_TOKENS))) - mask = pc.or_(pc.or_(is_null, blank), pc.fill_null(token, False)) - return np.asarray(pc.fill_null(mask, True)) + return _blank_mask(trimmed) | np.asarray(pc.fill_null(token, False)) def _regex_ok(values: pa.Array, pattern: str) -> np.ndarray: @@ -244,9 +254,12 @@ def _encode_annotation_column( strings = _as_string(arr) trimmed = pc.utf8_trim_whitespace(strings) - missing = _missing_mask(trimmed) + blank = _blank_mask(trimmed) # --- numeric inference (conversion.ts:71-125) --------------------------- # + # Only here does a MISSING_TOKENS spelling count as absent, so a column of + # ``NA`` stays categorical rather than turning into an all-NaN numeric. + missing = _missing_mask(trimmed) if not missing.all(): numeric_ok = _regex_ok(trimmed, JS_NUMBER_RE) | missing if numeric_ok.all(): @@ -263,11 +276,14 @@ def _encode_annotation_column( return entry, pa.array(values, type=pa.float64()), [] # --- categorical: split cells into hits --------------------------------- # - cells = pc.if_else(pa.array(~missing), trimmed, pa.scalar(None, pa.string())) + # ``_blank_mask``, not ``_missing_mask``: v3 is a container encoding and must + # hand back the label it was given, so ``none``/``NA``/``null`` stay ordinary + # categories here and the browser folds them into NA on read as it always has. + cells = pc.if_else(pa.array(~blank), trimmed, pa.scalar(None, pa.string())) hit_lists = pc.split_pattern(cells, ";") row_of_hit = np.asarray(pc.list_parent_indices(hit_lists)) hits = pc.utf8_trim_whitespace(pc.list_flatten(hit_lists)) - keep = ~_missing_mask(hits) + keep = ~_blank_mask(hits) if not keep.all(): hits = hits.filter(pa.array(keep)) row_of_hit = row_of_hit[keep] @@ -346,7 +362,7 @@ def _encode_annotation_column( if has_scores: counts = _counts_i32(hit_score_count, f"column '{name}' scores") payloads.append((f"score_count:{name}", counts.tobytes())) - payloads.append((f"scores:{name}", score_values.astype(" tuple[pa.Table, pa.Table, pa.Table]: * hits and cells are whitespace-trimmed, and empty or missing-valued hits are dropped (``"A;;B"`` comes back ``"A;B"``, ``" A |IDA"`` as ``"A|IDA"``); - * a missing cell -- null, blank or a ``MISSING_TOKENS`` spelling -- comes back - as ``""``; + * a missing cell -- null or blank -- comes back as ``""`` (a cell spelled + ``none``/``NA``/``null`` is an ordinary label and comes back unchanged); * labels are re-encoded canonically, so ``%3b`` comes back as ``%3B``; - * scores round-trip through float32 and are re-spelled shortest-first, so - ``"0.5700"`` comes back as ``"0.57"``; + * scores are re-spelled shortest-first, so ``"0.5700"`` comes back as + ``"0.57"``; * a numeric column comes back in its ``sourceType`` when that is restorable and otherwise as its canonical v2 spelling, so an all-integral column spells ``100``, never ``100.0``; diff --git a/apps/protspace/tests/test_bundle_v3_decode.py b/apps/protspace/tests/test_bundle_v3_decode.py index 1703b62e..ffaefbcc 100644 --- a/apps/protspace/tests/test_bundle_v3_decode.py +++ b/apps/protspace/tests/test_bundle_v3_decode.py @@ -271,9 +271,15 @@ def test_hits_and_cells_are_trimmed_and_empty_hits_collapse(): assert cells(table, "pfam") == ["A;B", "A;B", "A;B"] -def test_missing_spellings_all_come_back_as_the_empty_string(): +def test_only_null_and_blank_cells_come_back_as_the_empty_string(): + """``none``/``NA``/``null`` are labels, not missing values. + + ``protspace style`` resolves a legend entry by its literal cell value, so + collapsing these would delete a real category (1383 of the 1587 + ``phosphatase.predicted_transmembrane`` rows are literally ``none``). + """ table = annotations_table(col=["A", "NA", "n/a", "None", "__NA__", " "]) - assert cells(table, "col") == ["A", "", "", "", "", ""] + assert cells(table, "col") == ["A", "NA", "n/a", "None", "__NA__", ""] def test_null_cells_come_back_as_the_empty_string(): @@ -301,13 +307,24 @@ def test_an_unscored_hit_in_a_scored_column_keeps_no_suffix(): assert cells(table, "col") == ["PF1|0.5;PF2", "PF3"] -def test_scores_round_trip_through_float32(): +def test_scores_round_trip_through_float64(): table = annotations_table(col=["A|0.5700", "A|1", "A|0.1", "A|1e-10,2.5"]) - # 0.5700 loses its trailing zero (float32 has no such notion) and an integral + # 0.5700 loses its trailing zero (a float has no such notion) and an integral # score keeps the JavaScript spelling ``[1].join(',') === '1'``. assert cells(table, "col") == ["A|0.57", "A|1", "A|0.1", "A|1e-10,2.5"] +def test_e_values_survive_the_round_trip(): + """float32 scores would spell these ``0``, ``0``, ``inf`` and ``1.2345679e+08``.""" + table = annotations_table(col=["A|1e-200", "A|1e-300", "A|1e40", "A|123456789"]) + assert cells(table, "col") == [ + "A|1e-200", + "A|1e-300", + "A|1e+40", + "A|123456789", + ] + + def test_an_int_column_re_spells_its_cells_canonically(): table = annotations_table(col=["1", "2.0", "+3", "4e1"]) assert cells(table, "col") == ["1", "2", "3", "40"] diff --git a/apps/protspace/tests/test_bundle_v3_encode.py b/apps/protspace/tests/test_bundle_v3_encode.py index af0b9f71..608a1701 100644 --- a/apps/protspace/tests/test_bundle_v3_encode.py +++ b/apps/protspace/tests/test_bundle_v3_encode.py @@ -214,11 +214,36 @@ def test_code_order_is_frequency_then_first_occurrence(): assert read(parts[0]).column("col").to_pylist() == [3, 0, 1, 2, 0, 1, 2, 0] -def test_missing_cells_are_minus_one_and_not_a_category(): +def test_only_blank_cells_are_minus_one(): + """A cell literally spelled ``none`` is a category, not a missing value. + + v3 is a container encoding: collapsing the six ``MISSING_TOKENS`` spellings + would rewrite the data (``phosphatase.predicted_transmembrane`` is 1383 of + 1587 rows of literal ``none``, and ``protspace style`` keys on that label). + The browser still folds them into NA at read time, on v2 and v3 alike. + """ table = make_annotations(col=["A", "", "NA", "n/a", "None", "__NA__", " ", "A"]) parts = encode(table) - assert labels_of(payloads_of(parts[3]), "col") == ["A"] - assert read(parts[0]).column("col").to_pylist() == [0, -1, -1, -1, -1, -1, -1, 0] + assert labels_of(payloads_of(parts[3]), "col") == [ + "A", + "NA", + "n/a", + "None", + "__NA__", + ] + assert read(parts[0]).column("col").to_pylist() == [0, -1, 1, 2, 3, 4, -1, 0] + + +def test_missing_tokens_only_gate_numeric_inference(): + """``MISSING_TOKENS`` survives for exactly one job: keeping ``NA`` non-numeric.""" + parts = encode(make_annotations(col=["1", "NA", "none"])) + assert manifest_of(parts[0])["columns"]["col"]["kind"] == "numeric" + + # ...so a column made only of them cannot become an all-NaN numeric column, + # and lands in the categorical path with its spellings intact. + parts = encode(make_annotations(col=["NA", "none", "NA"])) + assert manifest_of(parts[0])["columns"]["col"]["kind"] == "categorical" + assert labels_of(payloads_of(parts[3]), "col") == ["NA", "none"] def test_scored_multi_column_csr_and_payloads(): @@ -230,11 +255,28 @@ def test_scored_multi_column_csr_and_payloads(): assert list(np.frombuffer(payloads["csr:pfam"], " Date: Sun, 6 Sep 2026 02:15:37 +0200 Subject: [PATCH 22/31] fix(bundle): reject corrupt v3 label lengths and CSR counts Also make the int64 spelling guard per value and drop the dead zero-chunk arm. --- .../src/protspace/data/io/bundle_v3.py | 59 ++++++-- apps/protspace/tests/test_bundle_v3_decode.py | 130 +++++++++++++++++- 2 files changed, 176 insertions(+), 13 deletions(-) diff --git a/apps/protspace/src/protspace/data/io/bundle_v3.py b/apps/protspace/src/protspace/data/io/bundle_v3.py index d5f3fe6b..abc43f01 100644 --- a/apps/protspace/src/protspace/data/io/bundle_v3.py +++ b/apps/protspace/src/protspace/data/io/bundle_v3.py @@ -579,13 +579,17 @@ def _read(part: bytes) -> pa.Table: def _flat(column: pa.ChunkedArray | pa.Array) -> pa.Array: - """One contiguous Arrow array (``ListArray.from_arrays`` refuses chunks).""" + """One contiguous Arrow array (``ListArray.from_arrays`` refuses chunks). + + Every v3 part is one row group, so the single-chunk branch is what actually + runs (a zero-row part still reads back as one empty chunk). The concat stays + because pyarrow splits a column past the 2 GB BinaryArray limit into several + chunks, and taking chunk 0 there would silently truncate the column. + """ if not isinstance(column, pa.ChunkedArray): return column if column.num_chunks == 1: return column.chunk(0) - if column.num_chunks == 0: - return pa.array([], type=column.type) return pa.concat_arrays(column.chunks) @@ -601,19 +605,44 @@ def _read_payloads(part: bytes) -> dict[str, bytes]: def _read_labels(payloads: dict[str, bytes], name: str) -> list[str]: - """Slice ``dict:`` by the prefix sum of its per-label byte lengths.""" + """Slice ``dict:`` by the prefix sum of its per-label byte lengths. + + A v3 bundle is user-supplied input and Python slicing clamps, so a corrupt + length array would silently yield duplicated and empty labels instead of an + error. The lengths must therefore tile the blob exactly. + """ blob = payloads[f"dict:{name}"] lengths = np.frombuffer(payloads[f"dict:{name}:len"], " pa.Array: - """Prefix-sum per-element ``counts`` into list offsets, then join each list.""" - offsets = np.concatenate(([0], np.cumsum(counts, dtype=np.int64))).astype(np.int32) +def _list_join( + counts: np.ndarray, values: pa.Array, separator: str, what: str +) -> pa.Array: + """Prefix-sum per-element ``counts`` into list offsets, then join each list. + + A total below ``len(values)`` silently empties the tail lists and a negative + count misaligns them, so a user-supplied bundle has to tile ``values`` + exactly. (A total above it already raises inside Arrow.) + """ + ends = np.cumsum(counts, dtype=np.int64) + total = int(ends[-1]) if ends.size else 0 + if bool((counts < 0).any()) or total != len(values): + raise ValueError( + f"{what} is corrupt: {counts.size} count(s) totalling {total} over " + f"{len(values)} value(s)" + ) + offsets = np.concatenate(([0], ends)).astype(np.int32) lists = pa.ListArray.from_arrays(pa.array(offsets, type=pa.int32()), values) return pc.binary_join(lists, separator) @@ -644,9 +673,15 @@ def _decode_numeric(column: pa.ChunkedArray, entry: dict[str, Any]) -> pa.Array: finite = np.where(present, values, 0.0) # ``str(2.0)`` is ``"2.0"`` but an int-typed v2 column spells it ``"2"``, and # numpy's float repr is Python's, so int columns take the int64 detour. The - # magnitude guard keeps a value past int64 out of an undefined cast. - if entry.get("numericType") == "int" and np.abs(finite).max(initial=0.0) < 2.0**63: - text = finite.astype(np.int64).astype(str) + # magnitude guard keeps a value past int64 out of an undefined cast, and it is + # per value: one 1e19 cell must not re-spell the whole column as floats. + if entry.get("numericType") == "int": + small = np.abs(finite) < 2.0**63 + text = np.where( + small, + np.where(small, finite, 0.0).astype(np.int64).astype(str), + finite.astype(str), + ) else: text = finite.astype(str) return pc.if_else( @@ -692,7 +727,7 @@ def _decode_multi( ) scored = pc.if_else( pa.array(per_hit > 0), - _list_join(per_hit, text, ","), + _list_join(per_hit, text, ",", f"payload 'score_count:{name}'"), pa.scalar(None, pa.string()), ) suffix = scored if suffix is None else pc.coalesce(suffix, scored) @@ -703,7 +738,7 @@ def _decode_multi( hits = pc.coalesce(pc.binary_join_element_wise(hits, suffix, "|"), hits) counts = _flat(column).to_numpy(zero_copy_only=False) - return _list_join(counts, hits, ";") + return _list_join(counts, hits, ";", f"column '{name}__count'") def _decode_projections( diff --git a/apps/protspace/tests/test_bundle_v3_decode.py b/apps/protspace/tests/test_bundle_v3_decode.py index ffaefbcc..4a7f014d 100644 --- a/apps/protspace/tests/test_bundle_v3_decode.py +++ b/apps/protspace/tests/test_bundle_v3_decode.py @@ -24,7 +24,7 @@ FORMAT_VERSION_KEY, stamp_format_version, ) -from protspace.data.io.bundle_v3 import MANIFEST_KEY, decode_v3, encode_v3 +from protspace.data.io.bundle_v3 import MANIFEST_KEY, _flat, decode_v3, encode_v3 from protspace.data.io.predictions import add_overlay_columns from protspace.data.processors.base_processor import BaseProcessor @@ -330,6 +330,35 @@ def test_an_int_column_re_spells_its_cells_canonically(): assert cells(table, "col") == ["1", "2", "3", "40"] +def test_an_int_past_the_int64_range_only_re_spells_itself(): + """The magnitude guard is per value: one huge cell is not a column-wide float.""" + table = annotations_table(col=["100", "250", "10000000000000000000"]) + assert cells(table, "col") == ["100", "250", "1e+19"] + + +def test_a_large_string_column_comes_back_as_its_v2_spelling(): + """``large_string`` is a parseable alias but not a numeric type to restore.""" + source = stamp_format_version( + pa.table( + { + "protein_id": ["p0", "p1"], + "length": pa.array(["100", "200"], type=pa.large_string()), + } + ) + ) + decoded = round_trip(source, (2,))[0] + assert ( + json.loads( + pq.read_table( + io.BytesIO(encode_v3(source, *projection_tables(2, (2,)))[0]) + ).schema.metadata[MANIFEST_KEY] + )["columns"]["length"]["sourceType"] + == "large_string" + ) + assert decoded.schema.field("length").type == pa.string() + assert decoded.column("length").to_pylist() == ["100", "200"] + + def test_a_bool_column_comes_back_as_the_python_spelling(): """``sourceType`` restoration is numeric-only; a bool stays v2's ``True``/``False``.""" source = stamp_format_version( @@ -381,6 +410,105 @@ def test_rejects_an_unknown_kind(): decode_v3(parts) +# --------------------------------------------------------------------------- # +# corrupt payloads (a v3 bundle is user-supplied input) +# --------------------------------------------------------------------------- # + + +def rewrite(part: bytes, edit) -> bytes: + """Read a part, hand the table to ``edit``, write the result back.""" + buffer = io.BytesIO() + pq.write_table(edit(pq.read_table(io.BytesIO(part))), buffer) + return buffer.getvalue() + + +def corrupt_payload(parts: list[bytes], name: str, data: bytes) -> list[bytes]: + """Replace one payload row of part 6.""" + + def edit(table): + names = table.column("name").to_pylist() + blobs = table.column("data").to_pylist() + blobs[names.index(name)] = data + return pa.table({"name": names, "data": blobs}) + + return [*parts[:3], rewrite(parts[3], edit)] + + +def encoded(**columns: list[str]) -> list[bytes]: + source = annotations_table(**columns) + return list(encode_v3(source, *projection_tables(source.num_rows, (2,)))) + + +def test_rejects_label_lengths_that_do_not_tile_the_blob(): + """Python slicing clamps, so this would yield duplicated and empty labels.""" + parts = corrupt_payload( + encoded(col=["Alpha", "Beta", "Gamma"]), + "dict:col:len", + np.array([100, 100, 100], dtype="2 GB column reads back chunked.""" + chunked = pa.chunked_array( + [pa.array([1, 2], type=pa.int32()), pa.array([3], type=pa.int32())] + ) + assert _flat(chunked).to_pylist() == [1, 2, 3] + assert isinstance(_flat(chunked), pa.Array) + + # --------------------------------------------------------------------------- # # the real shipped bundle # --------------------------------------------------------------------------- # From 4569a691659d13f5dee0821a85b2d2a64b83ae54 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 02:41:57 +0200 Subject: [PATCH 23/31] feat(bundle): emit v3 containers and decode them back at every read Six slots always, re-encode the core when annotations change, read_tables. --- .../src/protspace/data/io/__init__.py | 2 + .../protspace/src/protspace/data/io/bundle.py | 265 ++++++++----- apps/protspace/tests/test_bundle_overlay.py | 18 +- .../tests/test_bundle_v3_container.py | 350 ++++++++++++++++++ apps/protspace/tests/test_bundle_version.py | 9 +- apps/protspace/tests/test_stats_bundle.py | 40 +- apps/protspace/tests/test_transfer_cli.py | 64 ++-- 7 files changed, 617 insertions(+), 131 deletions(-) create mode 100644 apps/protspace/tests/test_bundle_v3_container.py diff --git a/apps/protspace/src/protspace/data/io/__init__.py b/apps/protspace/src/protspace/data/io/__init__.py index 0c0b0fee..1cf6124c 100644 --- a/apps/protspace/src/protspace/data/io/__init__.py +++ b/apps/protspace/src/protspace/data/io/__init__.py @@ -15,6 +15,7 @@ read_settings_from_bytes, read_settings_from_file, read_statistics_from_bundle, + read_tables, replace_settings_in_bundle, write_bundle, ) @@ -32,6 +33,7 @@ "extract_bundle_to_dir", "read_bundle", "read_statistics_from_bundle", + "read_tables", "write_bundle", "replace_settings_in_bundle", "create_settings_parquet", diff --git a/apps/protspace/src/protspace/data/io/bundle.py b/apps/protspace/src/protspace/data/io/bundle.py index 2c17504e..d86ee5a7 100644 --- a/apps/protspace/src/protspace/data/io/bundle.py +++ b/apps/protspace/src/protspace/data/io/bundle.py @@ -5,10 +5,20 @@ fourth part carries settings (annotation colours, shapes, etc.); an optional fifth part carries projection statistics. -Positional layout: ``core(3) + settings? + statistics?``. When statistics are -present but settings are absent, the fourth part is written as **zero bytes** so -the statistics part is unambiguously the fifth — readers and writers branch on -the fourth part's emptiness, not on the raw part count. +Positional layout: ``core(3) + settings? + statistics?`` for a legacy (v1/v2) +container, and always ``core(3) + settings + statistics + payloads`` for v3. +When statistics are present but settings are absent, the fourth part is written +as **zero bytes** so the statistics part is unambiguously the fifth — readers +and writers branch on the fourth part's emptiness, not on the raw part count. +A v3 container writes both of those slots unconditionally (zero bytes when +absent) because the browser reads its payload part positionally, from +``parts[5]``. + +v3 is a *container-boundary* encoding: :func:`write_bundle` takes the v2-shaped +tables the pipeline already builds and emits v3 parts, and every read here +(:func:`read_tables`, :func:`read_bundle`, :func:`extract_bundle_to_dir`) hands +back v2-shaped tables again, so nothing above this module has to know. See +:mod:`protspace.data.io.bundle_v3`. """ import io @@ -21,7 +31,8 @@ import pyarrow as pa import pyarrow.parquet as pq -from protspace.data.annotations.encoding import stamp_format_version +from protspace.data.annotations.encoding import FORMAT_VERSION_KEY, stamp_format_version +from protspace.data.io.bundle_v3 import CONTAINER_VERSION, decode_v3, encode_v3 logger = logging.getLogger(__name__) @@ -37,22 +48,48 @@ STATISTICS_FILENAME = "statistics.parquet" -def _parse_bundle(bundle_path: Path) -> tuple[list[bytes], bytes | None, bytes | None]: - """Read a bundle → ``(core_parts, settings_bytes, statistics_bytes)``. +def _part_container_version(part: bytes) -> int: + """The ``protspace_format_version`` in a part's parquet footer (1 if absent).""" + metadata = pq.read_metadata(io.BytesIO(part)).metadata or {} + try: + return int(metadata.get(FORMAT_VERSION_KEY, b"1")) + except (TypeError, ValueError): + return 1 + - The single place the on-disk layout is decoded: reads the file, validates the - 3-to-5 part count, and normalises the optional parts (the zero-byte settings - sentinel and an absent/empty statistics part both become ``None``). +def _split(data: bytes) -> tuple[list[bytes], bytes | None, bytes | None, bytes | None]: + """Split raw bundle bytes → ``(core_parts, settings, statistics, payloads)``. + + Six parts is v3 and the part-1 footer has to say so; three to five parts is a + legacy container, which has no payloads. The optional parts are normalised + (the zero-byte settings sentinel and an absent/empty statistics part both + become ``None``), so callers never branch on the raw part count. """ - with open(bundle_path, "rb") as f: - parts = f.read().split(PARQUET_BUNDLE_DELIMITER) + parts = data.split(PARQUET_BUNDLE_DELIMITER) + + if len(parts) < 3 or len(parts) > 6: + raise ValueError(f"Expected 3 to 6 parts in parquetbundle, found {len(parts)}") - if len(parts) < 3 or len(parts) > 5: - raise ValueError(f"Expected 3 to 5 parts in parquetbundle, found {len(parts)}") + payloads = None + if len(parts) == 6: + version = _part_container_version(parts[0]) + if version != CONTAINER_VERSION: + raise ValueError( + f"6-part parquetbundle declares container version {version}, " + f"expected {CONTAINER_VERSION}" + ) + payloads = parts[5] settings = parts[3] if len(parts) >= 4 and parts[3] else None - statistics = parts[4] if len(parts) == 5 and parts[4] else None - return parts[:3], settings, statistics + statistics = parts[4] if len(parts) >= 5 and parts[4] else None + return parts[:3], settings, statistics, payloads + + +def _parse_bundle( + bundle_path: Path, +) -> tuple[list[bytes], bytes | None, bytes | None, bytes | None]: + """:func:`_split` over a file: the single place the on-disk layout is decoded.""" + return _split(Path(bundle_path).read_bytes()) def _table_to_parquet_bytes(table: pa.Table) -> bytes: @@ -99,12 +136,70 @@ def _check_no_delimiter(part_bytes: bytes) -> None: ) +def _write_parts( + path: Path, + core: list[bytes], + settings: bytes | None = None, + statistics: bytes | None = None, + payloads: bytes | None = None, +) -> None: + """Assemble and atomically write one container from its already-serialized parts. + + The single writer for every bundle this module produces. A v3 container + (``payloads`` given) always emits **six** slots: the browser reads the + payloads from ``parts[5]`` positionally, so omitting an absent settings or + statistics part would file the payloads under statistics and the reader would + report no payloads part. A legacy container keeps the old trailing-optional + layout. Every part, part 6 included, is checked for the delimiter — a label + carrying those bytes would corrupt the split on read-back. + """ + if len(core) != 3: + raise ValueError(f"a parquetbundle needs exactly 3 core parts, got {len(core)}") + + if payloads is not None: + parts = [*core, settings or b"", statistics or b"", payloads] + else: + parts = list(core) + if settings is not None or statistics is not None: + parts.append(settings if settings is not None else b"") + if statistics is not None: + parts.append(statistics) + + for part in parts: + _check_no_delimiter(part) + + _atomic_write_bytes(path, PARQUET_BUNDLE_DELIMITER.join(parts)) + + +def read_tables( + path_or_bytes: Path | str | bytes, +) -> tuple[pa.Table, pa.Table, pa.Table]: + """Read a bundle's three core tables in their v2 shape. + + A v3 container is decoded (all-string annotation cells, long-format + projections, footer re-stamped ``protspace_format_version=2``); a legacy + container's parts are read as they are, so a v1 bundle stays v1-stamped and + is never silently migrated. + """ + data = ( + path_or_bytes + if isinstance(path_or_bytes, bytes) + else Path(path_or_bytes).read_bytes() + ) + core, _settings, _statistics, payloads = _split(data) + if payloads is not None: + return decode_v3([*core, payloads]) + annotations, metadata, projections = (pq.read_table(io.BytesIO(p)) for p in core) + return annotations, metadata, projections + + def extract_bundle_to_dir(bundle_path: Path, target_dir: Path | None = None) -> str: """Extract a .parquetbundle into separate parquet files on disk. - Supports bundles with 3 parts (core data only), 4 parts (core + settings), - or 5 parts (core + settings + statistics, where the settings part may be - zero bytes). + Supports legacy bundles with 3 parts (core data only), 4 parts (core + + settings) or 5 parts (core + settings + statistics, where the settings part + may be zero bytes), and 6-part v3 bundles — whose core is decoded back to the + v2 shape so everything downstream reads the files it always did. Args: bundle_path: Path to the .parquetbundle file. @@ -120,11 +215,17 @@ def extract_bundle_to_dir(bundle_path: Path, target_dir: Path | None = None) -> target_dir = Path(target_dir) target_dir.mkdir(parents=True, exist_ok=True) - core, settings, statistics = _parse_bundle(bundle_path) + core, settings, statistics, payloads = _parse_bundle(bundle_path) - for part_bytes, filename in zip(core, CORE_FILENAMES, strict=False): - if part_bytes: - (target_dir / filename).write_bytes(part_bytes) + if payloads is not None: + for table, filename in zip( + decode_v3([*core, payloads]), CORE_FILENAMES, strict=True + ): + pq.write_table(table, str(target_dir / filename)) + else: + for part_bytes, filename in zip(core, CORE_FILENAMES, strict=False): + if part_bytes: + (target_dir / filename).write_bytes(part_bytes) if settings: (target_dir / SETTINGS_FILENAME).write_bytes(settings) if statistics: @@ -134,24 +235,27 @@ def extract_bundle_to_dir(bundle_path: Path, target_dir: Path | None = None) -> def read_bundle(bundle_path: Path) -> tuple[list[bytes], dict | None]: - """Read a bundle and return raw core part bytes plus parsed settings. + """Read a bundle and return v2-shaped core part bytes plus parsed settings. The return shape is preserved (``(core_parts, settings)``) so existing callers keep working; use :func:`read_statistics_from_bundle` for the - optional fifth part. + optional statistics part. A v3 container is decoded and re-serialized, so + the parts callers ``pq.read_table`` are always the v2 shape — prefer + :func:`read_tables` when you want the tables themselves and not the bytes. Returns: (core_parts_bytes, settings_dict_or_None) """ - core, settings_bytes, _ = _parse_bundle(bundle_path) + core, settings_bytes, _statistics, payloads = _parse_bundle(bundle_path) + if payloads is not None: + core = [_table_to_parquet_bytes(t) for t in decode_v3([*core, payloads])] settings = read_settings_from_bytes(settings_bytes) if settings_bytes else None return core, settings def read_statistics_from_bundle(bundle_path: Path) -> bytes | None: """Return the raw statistics parquet bytes (fifth part), or None if absent.""" - _, _, statistics = _parse_bundle(bundle_path) - return statistics + return _parse_bundle(bundle_path)[2] def write_bundle( @@ -162,40 +266,36 @@ def write_bundle( ) -> None: """Write Arrow tables (and optional settings/statistics) to a .parquetbundle. + The tables come in v2-shaped (all-string annotation cells, long-format + projections) and go out as a six-part v3 container. + Args: tables: List of 3 Arrow tables (annotations, projections_metadata, projections_data). bundle_path: Output file path. settings: Optional settings dict to include as 4th part. statistics: Optional projection-statistics Arrow table to include as the - 5th part. When given without ``settings``, a zero-byte settings slot - is written so the statistics part stays at position five. + 5th part. A zero-byte slot is written for whichever of the two is + absent, so the v3 payloads part stays at position six. """ - buf = io.BytesIO() - for i, table in enumerate(tables): - if i > 0: - buf.write(PARQUET_BUNDLE_DELIMITER) - part_bytes = _table_to_parquet_bytes(table) - _check_no_delimiter(part_bytes) - buf.write(part_bytes) - - # A settings slot must exist whenever statistics follow it, so the parts - # keep fixed positions (settings = 4th, statistics = 5th). - if settings is not None or statistics is not None: - buf.write(PARQUET_BUNDLE_DELIMITER) - if settings is not None: - settings_bytes = create_settings_parquet(settings) - _check_no_delimiter(settings_bytes) - buf.write(settings_bytes) - # else: zero-byte settings slot keeps statistics at position five - - if statistics is not None: - buf.write(PARQUET_BUNDLE_DELIMITER) - stats_bytes = _table_to_parquet_bytes(statistics) - _check_no_delimiter(stats_bytes) - buf.write(stats_bytes) - - _atomic_write_bytes(bundle_path, buf.getvalue()) + if len(tables) != 3: + raise ValueError( + f"write_bundle expects 3 core tables (annotations, projections_metadata, " + f"projections_data), got {len(tables)}" + ) + + annotations, projections_metadata, projections_data = tables + part1, part2, part3, payloads = encode_v3( + annotations, projections_metadata, projections_data + ) + + _write_parts( + bundle_path, + [part1, part2, part3], + create_settings_parquet(settings) if settings is not None else None, + _table_to_parquet_bytes(statistics) if statistics is not None else None, + payloads, + ) logger.info(f"Saved bundled output to: {bundle_path}") @@ -206,20 +306,14 @@ def replace_settings_in_bundle( ) -> None: """Append or replace the settings (4th) part in a bundle. - The three core parts are preserved byte-for-byte, and an existing statistics - (5th) part is preserved so styling a statistics-bearing bundle is non-lossy. + Every other part is preserved byte-for-byte, so a legacy bundle stays legacy + and a v3 bundle keeps its payloads; an existing statistics part survives, so + styling a statistics-bearing bundle is non-lossy. """ - core, _, statistics = _parse_bundle(input_path) - - # core(3) + new settings, preserving a trailing statistics part if present. - settings_bytes = create_settings_parquet(settings) - _check_no_delimiter(settings_bytes) - new_parts = [*core, settings_bytes] - if statistics is not None: - new_parts.append(statistics) - new_content = PARQUET_BUNDLE_DELIMITER.join(new_parts) - - _atomic_write_bytes(output_path, new_content) + core, _settings, statistics, payloads = _parse_bundle(input_path) + _write_parts( + output_path, core, create_settings_parquet(settings), statistics, payloads + ) def replace_annotations_in_bundle( @@ -229,33 +323,30 @@ def replace_annotations_in_bundle( ) -> None: """Replace the annotations (1st) part of a bundle, preserving the rest. - Projection parts (2nd, 3rd) are kept byte-for-byte; existing settings (4th) - and statistics (5th) parts are carried over unchanged. + The whole v3 core is re-encoded, not just part 1: the payloads part holds the + label dictionaries and CSR buffers *for* part 1, so keeping the old one next + to new annotations would leave stale payloads behind. Settings and + statistics are carried over unchanged. A legacy input container comes out as + v3, which is correct — this is a write, and every write emits v3. """ - core, settings, statistics = _parse_bundle(input_path) + data = Path(input_path).read_bytes() + _core, settings, statistics, _payloads = _split(data) # Re-stamp the format version at this single annotations-write chokepoint. # pyarrow table ops (rename_columns, concat) drop schema metadata, and # callers (transfer, prediction overlay) build the replacement table from - # exactly such ops — so without this the stamp is silently lost and a v2 - # bundle re-reads as v1 (raw %XX names). Callers must provide v2-safe cells; - # transfer explicitly migrates legacy v1 categorical grammar before it - # reaches this write boundary. + # exactly such ops — so without this the stamp is silently lost and the + # encoder would migrate an already-v2 table a second time, double-escaping + # every reserved character. Callers must provide v2-safe cells; transfer + # explicitly migrates legacy v1 categorical grammar before this boundary. annotations_table = stamp_format_version(annotations_table) - new_annotations_bytes = _table_to_parquet_bytes(annotations_table) - _check_no_delimiter(new_annotations_bytes) - - # Preserve the projection parts byte-for-byte; keep the settings/statistics - # tail with the same zero-byte-settings sentinel write_bundle uses, so a - # statistics-bearing bundle round-trips without losing its 5th part. - new_parts = [new_annotations_bytes, core[1], core[2]] - if settings is not None or statistics is not None: - new_parts.append(settings if settings is not None else b"") - if statistics is not None: - new_parts.append(statistics) + _annotations, projections_metadata, projections_data = read_tables(data) + part1, part2, part3, payloads = encode_v3( + annotations_table, projections_metadata, projections_data + ) - _atomic_write_bytes(output_path, PARQUET_BUNDLE_DELIMITER.join(new_parts)) + _write_parts(output_path, [part1, part2, part3], settings, statistics, payloads) logger.info(f"Wrote bundle with updated annotations to: {output_path}") diff --git a/apps/protspace/tests/test_bundle_overlay.py b/apps/protspace/tests/test_bundle_overlay.py index 0b90bc6f..54a456e9 100644 --- a/apps/protspace/tests/test_bundle_overlay.py +++ b/apps/protspace/tests/test_bundle_overlay.py @@ -16,8 +16,15 @@ def _tables(): annotations = pa.table({"identifier": ["A", "B"], "cat": ["x", "y"]}) - proj_meta = pa.table({"name": ["PCA 2"], "dims": [2]}) - proj_data = pa.table({"id": ["A", "B"], "x": [0.0, 1.0], "y": [0.0, 1.0]}) + proj_meta = pa.table({"projection_name": ["PCA 2"], "dimensions": [2]}) + proj_data = pa.table( + { + "projection_name": ["PCA 2", "PCA 2"], + "identifier": ["A", "B"], + "x": [0.0, 1.0], + "y": [0.0, 1.0], + } + ) return [annotations, proj_meta, proj_data] @@ -37,12 +44,15 @@ def test_replaces_annotations_keeps_other_parts(tmp_path): parts, settings = read_bundle(out) assert "cat__pred_value" in _read_part(parts[0]).column_names - # Projections preserved byte-for-byte. - assert _read_part(parts[1]).column_names == ["name", "dims"] + # Projections preserved. + assert _read_part(parts[1]).column_names == ["projection_name", "dimensions"] assert _read_part(parts[2]).to_pydict()["x"] == [0.0, 1.0] def test_projection_parts_preserved_byte_for_byte(tmp_path): + """A v3 rewrite re-encodes the whole core (part 6 holds part 1's payloads), + so the projection parts are re-derived rather than copied — and still come + out byte-identical, which is what pins the encoder as deterministic.""" src = tmp_path / "in.parquetbundle" out = tmp_path / "out.parquetbundle" write_bundle(_tables(), src, settings={"foo": 1}) diff --git a/apps/protspace/tests/test_bundle_v3_container.py b/apps/protspace/tests/test_bundle_v3_container.py new file mode 100644 index 00000000..ed4f7ddb --- /dev/null +++ b/apps/protspace/tests/test_bundle_v3_container.py @@ -0,0 +1,350 @@ +"""The v3 container boundary (``data/io/bundle``). + +``bundle_v3`` owns the codec; this file owns the *container* around it: every +write emits a six-part v3 bundle, every read hands v2-shaped tables back, and a +legacy (v1/v2) bundle is still read exactly as it was written rather than +silently migrated. + +The three things that would break the browser if they regressed: + +* part 6 is positionally pinned (the reader takes payloads from ``parts[5]``), + so a v3 container always writes the settings and statistics slots, zero bytes + when absent; +* part 6 holds the label dictionaries *for* part 1, so any write that changes + the annotations has to re-encode both; +* the delimiter guard has to cover part 6, where the labels now live. +""" + +import io +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from protspace.data.annotations.encoding import ( + FORMAT_VERSION_KEY, + read_format_version, + stamp_format_version, +) +from protspace.data.io.bundle import ( + PARQUET_BUNDLE_DELIMITER, + _write_parts, + extract_bundle_to_dir, + read_bundle, + read_tables, + replace_annotations_in_bundle, + replace_settings_in_bundle, + write_bundle, +) +from tests.test_bundle_v3_decode import annotations_table, projection_tables + + +def pipeline_tables(dimensions=(2, 3)): + """The three v2-shaped tables the prepare pipeline hands ``write_bundle``.""" + annotations = annotations_table( + kingdom=["Bacteria", "Archaea", "Bacteria"], + pfam=["PF00001 (7tm%3B1)|1e-10,2.5", "", "PF00002|0.5;PF00003"], + go_mf=["GO:0005524|IDA", "GO:0016787|ECO:0000269", ""], + length=["120", "", "340"], + ) + metadata, data = projection_tables(annotations.num_rows, dimensions) + return [annotations, metadata, data] + + +def legacy_bundle(path: Path, *, stamp: bool = True, settings: bytes | None = None): + """Write a pre-v3 container by hand and return its raw parts. + + ``write_bundle`` cannot build one any more, and that is the point: a legacy + bundle has to keep reading back byte-for-byte as it was written. + """ + annotations = pa.table( + {"protein_id": ["p0", "p1"], "cat": ["ACC (Name%3B part)|EXP", "plain"]} + ) + if stamp: + annotations = stamp_format_version(annotations) + metadata = pa.table({"projection_name": ["PCA 2"], "dimensions": [2]}) + data = pa.table( + { + "projection_name": ["PCA 2", "PCA 2"], + "identifier": ["p0", "p1"], + "x": [0.0, 1.0], + "y": [2.0, 3.0], + } + ) + + def serialized(table): + buf = io.BytesIO() + pq.write_table(table, buf) + return buf.getvalue() + + parts = [serialized(t) for t in (annotations, metadata, data)] + if settings is not None: + parts.append(settings) + path.write_bytes(PARQUET_BUNDLE_DELIMITER.join(parts)) + return parts + + +def parts_of(path: Path) -> list[bytes]: + return path.read_bytes().split(PARQUET_BUNDLE_DELIMITER) + + +def read(part: bytes) -> pa.Table: + return pq.read_table(io.BytesIO(part)) + + +# --------------------------------------------------------------------------- # +# layout +# --------------------------------------------------------------------------- # + + +def test_write_bundle_always_emits_six_parts(tmp_path): + """Part 6 is read positionally, so the settings and statistics slots exist + even when empty; a five-slot v3 bundle would file the payloads under + statistics and the browser would report no payloads part.""" + path = tmp_path / "b.parquetbundle" + write_bundle(pipeline_tables(), path) + + parts = parts_of(path) + assert len(parts) == 6 + assert parts[3] == b"" and parts[4] == b"" + assert read(parts[0]).schema.metadata[FORMAT_VERSION_KEY] == b"3" + assert read(parts[5]).column_names == ["name", "data"] + + +def test_settings_and_statistics_keep_payloads_at_position_six(tmp_path): + path = tmp_path / "b.parquetbundle" + write_bundle( + pipeline_tables(), + path, + settings={"k": 1}, + statistics=pa.table({"metric": ["silhouette"], "value": [0.5]}), + ) + + parts = parts_of(path) + assert len(parts) == 6 + assert parts[3] and parts[4] + assert read(parts[5]).column_names == ["name", "data"] + + +def test_delimiter_in_a_label_is_caught_in_part_six(tmp_path): + """v3 moves labels out of part 1 and into the part 6 dictionary blob, so the + guard has to run there or a label carrying the reserved bytes corrupts the + split on read-back.""" + label = "x" + PARQUET_BUNDLE_DELIMITER.decode() + "y" + annotations, metadata, data = pipeline_tables() + annotations = annotations.set_column( + annotations.column_names.index("kingdom"), + "kingdom", + pa.array([label, "Archaea", "Bacteria"]), + ) + + with pytest.raises(ValueError, match="bundle delimiter"): + write_bundle( + [stamp_format_version(annotations), metadata, data], + tmp_path / "b.parquetbundle", + ) + + +def test_write_parts_checks_the_payloads_slot(tmp_path): + """The guard above only bites because ``_write_parts`` checks all six parts.""" + with pytest.raises(ValueError, match="bundle delimiter"): + _write_parts( + tmp_path / "b.parquetbundle", + [b"a", b"b", b"c"], + payloads=b"pay" + PARQUET_BUNDLE_DELIMITER + b"load", + ) + + +def test_six_parts_with_a_non_v3_footer_is_rejected(tmp_path): + """Six parts means v3; the part-1 footer has to agree, or the file is not a + bundle this reader understands.""" + path = tmp_path / "b.parquetbundle" + parts = legacy_bundle(tmp_path / "legacy.parquetbundle") + path.write_bytes(PARQUET_BUNDLE_DELIMITER.join([*parts, b"", b"", b"payloads"])) + + with pytest.raises(ValueError, match="container version 2"): + read_tables(path) + + +def test_too_many_parts_is_rejected(tmp_path): + path = tmp_path / "b.parquetbundle" + path.write_bytes(PARQUET_BUNDLE_DELIMITER.join([b""] * 7)) + + with pytest.raises(ValueError, match="Expected 3 to 6 parts"): + read_tables(path) + + +# --------------------------------------------------------------------------- # +# read_tables +# --------------------------------------------------------------------------- # + + +def test_write_then_read_tables_round_trips_pipeline_tables(tmp_path): + tables = pipeline_tables() + path = tmp_path / "b.parquetbundle" + write_bundle(tables, path) + + annotations, metadata, data = read_tables(path) + assert annotations.equals(tables[0]) + assert metadata.equals(tables[1]) + # Cell-for-cell, but not type-for-type: an all-null `z` leaves the pipeline + # as float64 and comes back float32 (a documented decode_v3 non-identity). + assert data.to_pydict() == tables[2].to_pydict() + assert data.schema.field("z").type == pa.float32() + # What comes back is the v2 cell grammar every Python consumer parses. + assert read_format_version(annotations) == 2 + + +def test_read_tables_accepts_raw_bytes(tmp_path): + tables = pipeline_tables() + path = tmp_path / "b.parquetbundle" + write_bundle(tables, path) + + assert read_tables(path.read_bytes())[0].equals(tables[0]) + + +def test_legacy_bundle_reads_back_unchanged(tmp_path): + path = tmp_path / "legacy.parquetbundle" + parts = legacy_bundle(path) + + annotations, metadata, data = read_tables(path) + assert annotations.equals(read(parts[0])) + assert metadata.equals(read(parts[1])) + assert data.equals(read(parts[2])) + assert read_format_version(annotations) == 2 + + +def test_legacy_v1_bundle_is_not_migrated_by_a_read(tmp_path): + """Reading must not upgrade: a v1 cell keeps its raw ``%XX`` spelling and its + missing stamp, so a consumer that gates decoding on the version still sees + v1 (double-escaping it here would be unrecoverable).""" + path = tmp_path / "legacy.parquetbundle" + legacy_bundle(path, stamp=False) + + annotations, _metadata, _data = read_tables(path) + assert read_format_version(annotations) == 1 + assert annotations.column("cat").to_pylist() == [ + "ACC (Name%3B part)|EXP", + "plain", + ] + + +# --------------------------------------------------------------------------- # +# rewrites +# --------------------------------------------------------------------------- # + + +def test_replace_settings_keeps_a_legacy_bundle_legacy(tmp_path): + src = tmp_path / "legacy.parquetbundle" + out = tmp_path / "styled.parquetbundle" + parts = legacy_bundle(src) + + replace_settings_in_bundle(src, out, {"new": 2}) + + out_parts = parts_of(out) + assert len(out_parts) == 4 # no payload slot invented + assert out_parts[:3] == parts[:3] # core preserved byte-for-byte + assert read_bundle(out)[1] == {"new": 2} + + +def test_replace_settings_keeps_the_payloads_of_a_v3_bundle(tmp_path): + src = tmp_path / "b.parquetbundle" + out = tmp_path / "styled.parquetbundle" + write_bundle(pipeline_tables(), src, settings={"old": 1}) + + replace_settings_in_bundle(src, out, {"new": 2}) + + in_parts, out_parts = parts_of(src), parts_of(out) + assert len(out_parts) == 6 + assert out_parts[5] == in_parts[5] + assert read_bundle(out)[1] == {"new": 2} + assert read_tables(out)[0].equals(pipeline_tables()[0]) + + +def test_replace_annotations_re_encodes_the_payloads(tmp_path): + """Part 6 holds part 1's label dictionaries, so rewriting part 1 while + carrying the old part 6 over would leave the codes pointing at stale labels. + Every label here is replaced, so a stale payload decodes to the old ones.""" + src = tmp_path / "b.parquetbundle" + out = tmp_path / "out.parquetbundle" + tables = pipeline_tables() + write_bundle(tables, src) + + replacement = annotations_table( + kingdom=["Fungi", "Viridiplantae", "Fungi"], + pfam=["PF09999|0.75", "", "PF08888"], + go_mf=["GO:0000001|IEA", "", ""], + length=["1", "2", "3"], + ) + replace_annotations_in_bundle(src, out, replacement) + + annotations, metadata, data = read_tables(out) + assert annotations.equals(replacement) + # The projections ride along untouched. + assert metadata.equals(tables[1]) + assert data.to_pydict() == tables[2].to_pydict() + + payload_blob = b"".join( + read(parts_of(out)[5]).column("data").to_pylist(), + ) + assert b"Fungi" in payload_blob + assert b"Bacteria" not in payload_blob # no stale dictionary left behind + + +def test_replace_annotations_upgrades_a_legacy_bundle(tmp_path): + """A rewrite is a write, and every write emits v3.""" + src = tmp_path / "legacy.parquetbundle" + out = tmp_path / "out.parquetbundle" + legacy_bundle(src) + + replacement = stamp_format_version( + pa.table({"protein_id": ["p0", "p1"], "cat": ["alpha", "beta"]}) + ) + replace_annotations_in_bundle(src, out, replacement) + + parts = parts_of(out) + assert len(parts) == 6 + assert read(parts[0]).schema.metadata[FORMAT_VERSION_KEY] == b"3" + annotations, _metadata, data = read_tables(out) + assert annotations.column("cat").to_pylist() == ["alpha", "beta"] + assert data.column("x").to_pylist() == [0.0, 1.0] + + +# --------------------------------------------------------------------------- # +# extraction +# --------------------------------------------------------------------------- # + + +def test_extract_v3_bundle_writes_v2_shaped_files(tmp_path): + src = tmp_path / "b.parquetbundle" + tables = pipeline_tables() + write_bundle(tables, src, settings={"k": 1}) + + out_dir = Path(extract_bundle_to_dir(src, tmp_path / "out")) + + extracted = [ + pq.read_table(str(out_dir / name)) + for name in ( + "selected_annotations.parquet", + "projections_metadata.parquet", + "projections_data.parquet", + ) + ] + assert extracted[0].equals(tables[0]) + assert extracted[1].equals(tables[1]) + assert extracted[2].to_pydict() == tables[2].to_pydict() + assert extracted[0].schema.metadata[FORMAT_VERSION_KEY] == b"2" + assert (out_dir / "settings.parquet").exists() + + +def test_extract_legacy_bundle_writes_the_raw_parts(tmp_path): + src = tmp_path / "legacy.parquetbundle" + parts = legacy_bundle(src, settings=b"") + + out_dir = Path(extract_bundle_to_dir(src, tmp_path / "out")) + + assert (out_dir / "selected_annotations.parquet").read_bytes() == parts[0] + assert (out_dir / "projections_metadata.parquet").read_bytes() == parts[1] + assert (out_dir / "projections_data.parquet").read_bytes() == parts[2] + assert not (out_dir / "settings.parquet").exists() # zero-byte slot diff --git a/apps/protspace/tests/test_bundle_version.py b/apps/protspace/tests/test_bundle_version.py index 5d6405c7..5e2a785a 100644 --- a/apps/protspace/tests/test_bundle_version.py +++ b/apps/protspace/tests/test_bundle_version.py @@ -94,8 +94,15 @@ def test_cli_bundle_command_stamps_format_version(tmp_path): pa.Table.from_pandas(metadata_df), proj_dir / "projections_metadata.parquet" ) + # The long layout `protspace project` actually writes (one row per protein + # per projection); `bundle` hands this table straight to write_bundle. data_df = pd.DataFrame( - {"identifier": ["P1", "P2"], "PCA_2_1": [0.1, 0.2], "PCA_2_2": [0.3, 0.4]} + { + "projection_name": ["PCA_2", "PCA_2"], + "identifier": ["P1", "P2"], + "x": [0.1, 0.2], + "y": [0.3, 0.4], + } ) pq.write_table(pa.Table.from_pandas(data_df), proj_dir / "projections_data.parquet") diff --git a/apps/protspace/tests/test_stats_bundle.py b/apps/protspace/tests/test_stats_bundle.py index d269b446..708b4828 100644 --- a/apps/protspace/tests/test_stats_bundle.py +++ b/apps/protspace/tests/test_stats_bundle.py @@ -18,8 +18,15 @@ def _core() -> list[pa.Table]: return [ pa.table({"protein_id": ["a", "b"]}), - pa.table({"projection_name": ["PCA_2"]}), - pa.table({"projection_name": ["PCA_2", "PCA_2"], "identifier": ["a", "b"]}), + pa.table({"projection_name": ["PCA_2"], "dimensions": [2]}), + pa.table( + { + "projection_name": ["PCA_2", "PCA_2"], + "identifier": ["a", "b"], + "x": [0.0, 1.0], + "y": [0.0, 1.0], + } + ), ] @@ -27,32 +34,39 @@ def _stats() -> pa.Table: return pa.table({"space_name": ["PCA_2"], "metric": ["silhouette"], "value": [0.5]}) -def _ndelims(path) -> int: - return path.read_bytes().count(PARQUET_BUNDLE_DELIMITER) +def _parts(path) -> list[bytes]: + return path.read_bytes().split(PARQUET_BUNDLE_DELIMITER) -def test_three_part_bundle_roundtrips(tmp_path): +def test_bundle_without_settings_or_stats_roundtrips(tmp_path): + """Every write is v3, so the container always has six slots — the settings + and statistics ones are zero bytes here, which is what keeps the payloads + part at position six where the browser reads it.""" p = tmp_path / "b.parquetbundle" write_bundle(_core(), p) - assert _ndelims(p) == 2 + parts = _parts(p) + assert len(parts) == 6 + assert parts[3] == b"" and parts[4] == b"" and parts[5] core, settings = read_bundle(p) assert len(core) == 3 and settings is None assert read_statistics_from_bundle(p) is None -def test_four_part_settings_only(tmp_path): +def test_settings_only(tmp_path): p = tmp_path / "b.parquetbundle" write_bundle(_core(), p, settings={"hello": "world"}) - assert _ndelims(p) == 3 + parts = _parts(p) + assert len(parts) == 6 and parts[3] and parts[4] == b"" _, settings = read_bundle(p) assert settings == {"hello": "world"} assert read_statistics_from_bundle(p) is None -def test_five_part_settings_and_stats(tmp_path): +def test_settings_and_stats(tmp_path): p = tmp_path / "b.parquetbundle" write_bundle(_core(), p, settings={"k": 1}, statistics=_stats()) - assert _ndelims(p) == 4 + parts = _parts(p) + assert len(parts) == 6 and parts[3] and parts[4] _, settings = read_bundle(p) assert settings == {"k": 1} stats_bytes = read_statistics_from_bundle(p) @@ -61,10 +75,12 @@ def test_five_part_settings_and_stats(tmp_path): assert table.column("metric")[0].as_py() == "silhouette" -def test_five_part_stats_only_empty_settings(tmp_path): +def test_stats_only_empty_settings(tmp_path): p = tmp_path / "b.parquetbundle" write_bundle(_core(), p, statistics=_stats()) - assert _ndelims(p) == 4 # zero-byte settings slot keeps stats at position 5 + parts = _parts(p) + # zero-byte settings slot keeps stats at position 5 and payloads at 6 + assert len(parts) == 6 and parts[3] == b"" and parts[4] core, settings = read_bundle(p) assert len(core) == 3 and settings is None assert read_statistics_from_bundle(p) is not None diff --git a/apps/protspace/tests/test_transfer_cli.py b/apps/protspace/tests/test_transfer_cli.py index c2b6156e..4950ddce 100644 --- a/apps/protspace/tests/test_transfer_cli.py +++ b/apps/protspace/tests/test_transfer_cli.py @@ -26,6 +26,21 @@ def _three_protein_inputs(extra_columns=None): return annotations, embeddings +def _projection_tables(identifiers, name="PCA 2"): + """The long-format projection tables the pipeline writes, for ``identifiers``.""" + return ( + pa.table({"projection_name": [name], "dimensions": [2]}), + pa.table( + { + "projection_name": [name] * len(identifiers), + "identifier": list(identifiers), + "x": [float(i) for i in range(len(identifiers))], + "y": [0.0] * len(identifiers), + } + ), + ) + + def _write_bundle_and_h5(tmp_path, *, id_col="protein_id", extra_columns=None): import h5py @@ -35,10 +50,7 @@ def _write_bundle_and_h5(tmp_path, *, id_col="protein_id", extra_columns=None): if extra_columns: cols.update(extra_columns) annotations = pa.table(cols) - proj_meta = pa.table({"name": ["PCA 2"], "dims": [2]}) - proj_data = pa.table( - {"id": ["TRINITY_1", "P00001"], "x": [0.0, 9.0], "y": [0.0, 0.0]} - ) + proj_meta, proj_data = _projection_tables(["TRINITY_1", "P00001"]) bundle_path = tmp_path / "in.parquetbundle" write_bundle([annotations, proj_meta, proj_data], bundle_path) @@ -365,10 +377,7 @@ def test_cli_end_to_end_protein_id_bundle(tmp_path): annotations = pa.table( {"protein_id": ["TRINITY_1", "P00001"], "protein_category": ["", "neurotoxin"]} ) - proj_meta = pa.table({"name": ["PCA 2"], "dims": [2]}) - proj_data = pa.table( - {"id": ["TRINITY_1", "P00001"], "x": [0.0, 9.0], "y": [0.0, 0.0]} - ) + proj_meta, proj_data = _projection_tables(["TRINITY_1", "P00001"]) bundle_path = tmp_path / "in.parquetbundle" write_bundle([annotations, proj_meta, proj_data], bundle_path) @@ -418,13 +427,12 @@ def test_cli_migrates_legacy_cells_and_encodes_reserved_source_id(tmp_path): from typer.testing import CliRunner from protspace.cli.app import app - from protspace.data.io.bundle import ( - PARQUET_BUNDLE_DELIMITER, - read_bundle, - write_bundle, - ) + from protspace.data.io.bundle import PARQUET_BUNDLE_DELIMITER, read_bundle source_id = "P0|ref;literal%3B" + # A genuine v1 container, assembled here rather than via write_bundle: every + # write now emits v3, whose annotations part comes back stamped v2, so a + # bundle built that way could not stand in for a legacy one. annotations = pa.table( { "protein_id": ["TRINITY_1", source_id], @@ -432,21 +440,18 @@ def test_cli_migrates_legacy_cells_and_encodes_reserved_source_id(tmp_path): "literal_percent": ["name%3Bpart", "plain"], } ) - proj_meta = pa.table({"name": ["PCA 2"], "dims": [2]}) - proj_data = pa.table( - {"id": ["TRINITY_1", source_id], "x": [0.0, 9.0], "y": [0.0, 0.0]} - ) - stamped_path = tmp_path / "stamped.parquetbundle" - write_bundle([annotations, proj_meta, proj_data], stamped_path) - parts, _ = read_bundle(stamped_path) - legacy_annotations = pq.read_table(io.BytesIO(parts[0])).replace_schema_metadata( - None - ) - first_part = io.BytesIO() - pq.write_table(legacy_annotations, first_part) + proj_meta, proj_data = _projection_tables(["TRINITY_1", source_id]) + + def _part(table): + buf = io.BytesIO() + pq.write_table(table, buf) + return buf.getvalue() + bundle_path = tmp_path / "legacy.parquetbundle" bundle_path.write_bytes( - PARQUET_BUNDLE_DELIMITER.join([first_part.getvalue(), parts[1], parts[2]]) + PARQUET_BUNDLE_DELIMITER.join( + [_part(annotations), _part(proj_meta), _part(proj_data)] + ) ) h5_path = tmp_path / "legacy.h5" @@ -568,4 +573,9 @@ def test_cli_transfer_without_rules_fills_missing_values(tmp_path): parts, _ = read_bundle(out_path) rows = {r["protein_id"]: r for r in pq.read_table(io.BytesIO(parts[0])).to_pylist()} assert rows["TRINITY_1"]["protein_category__pred_value"] == "neurotoxin" - assert rows["P00001"]["protein_category__pred_value"] is None + # A reference protein gets no prediction. The overlay writes null; a v3 + # container stores "absent" as a -1 dictionary code and spells it back as "" + # (a documented decode_v3 non-identity). Both readers of this column treat + # null and "" identically -- the browser's readCategoricalStorageValue and + # Python's to_display_value -- so the distinction is not observable. + assert rows["P00001"]["protein_category__pred_value"] == "" From 99b8d1c7b28fdb033ef2af05be74e9fb6d50b444 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 02:43:14 +0200 Subject: [PATCH 24/31] fix(core): read v3 scores as float64 and fold missing-value labels into NA Widen CsrScores.values to match the wire format. --- .../data-loader/utils/bundle-v3.test.ts | 85 +++++++++- .../components/data-loader/utils/bundle-v3.ts | 148 ++++++++++++++++-- packages/utils/src/types.ts | 7 +- .../src/visualization/eat-overlay.test.ts | 6 +- .../utils/src/visualization/eat-overlay.ts | 2 +- .../visualization/plot-data-accessors.test.ts | 2 +- .../slice-visualization-data.test.ts | 2 +- .../visualization/slice-visualization-data.ts | 2 +- 8 files changed, 227 insertions(+), 27 deletions(-) diff --git a/packages/core/src/components/data-loader/utils/bundle-v3.test.ts b/packages/core/src/components/data-loader/utils/bundle-v3.test.ts index 7953af05..f6e5d300 100644 --- a/packages/core/src/components/data-loader/utils/bundle-v3.test.ts +++ b/packages/core/src/components/data-loader/utils/bundle-v3.test.ts @@ -9,6 +9,7 @@ import { getProteinEvidence, getProteinScores, isCsrAnnotationData, + NA_DEFAULT_COLOR, NA_VALUE, type CsrAnnotationData, type VisualizationData, @@ -36,6 +37,7 @@ const enc = new TextEncoder(); const utf8 = (text: string) => enc.encode(text); const i32 = (...values: number[]) => new Uint8Array(new Int32Array(values).buffer); const f32 = (...values: number[]) => new Uint8Array(new Float32Array(values).buffer); +const f64 = (...values: number[]) => new Uint8Array(new Float64Array(values).buffer); type Column = { name: string; data: unknown[] | Int32Array | Float64Array | Float32Array }; @@ -126,7 +128,7 @@ const PAYLOADS: Record = { // Hit 3 is the first hit of P5, immediately after the empty interior row P4: its // score is what an off-by-one in the inserted-NA `hitEnd` would steal. 'score_count:go_bp': i32(2, 0, 1, 1, 0, 0, 0, 3, 0), - 'scores:go_bp': f32(1.5, 2.5, 9.75, 4, 0.5, 0.25, 0.125), + 'scores:go_bp': f64(1.5, 2.5, 9.75, 4, 0.5, 0.25, 0.125), 'evidence:go_bp': i32(-1, 0, 1, -1, -1, -1, 0, -1, -1), 'dict:__evidence': utf8('IDAECO:0000269'), 'dict:__evidence:len': i32(3, 11), @@ -149,7 +151,7 @@ const v3Bundle = (overrides: Record = {}) => * Every typed array `collectTransferables` names a buffer for, in a stable order, so a * dataset and its structured clone can be compared element by element. */ -const bulkViews = (data: VisualizationData): (Int32Array | Float32Array)[] => [ +const bulkViews = (data: VisualizationData): (Int32Array | Float32Array | Float64Array)[] => [ ...data.projections.map((projection) => projection.data as Float32Array), ...Object.values(data.annotation_data).flatMap((value) => value instanceof Int32Array @@ -247,6 +249,74 @@ describe('parquetbundle format v3', () => { expect(getProteinEvidence(data, 7, 'go_bp')).toEqual([null]); }); + it('keeps an E-value score exact, which float32 cannot', async () => { + // 1e-200 flushes to 0 and 1e40 saturates to Infinity in float32, and E-values are + // the canonical Pfam / InterPro score — so this is the whole reason the payload is + // float64. Reading it as float32 also halves the element count, which trips the + // score-count check first. + const payloads = { + ...PAYLOADS, + 'score_count:go_bp': i32(2, 0, 0, 0, 0, 0, 0, 0, 0), + 'scores:go_bp': f64(1e-200, 1e40), + }; + const { data } = await decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) })); + + expect(getProteinScores(data, 1, 'go_bp')).toEqual([[1e-200, 1e40], null]); + }); + + it('folds every missing-value spelling in a dictionary into ONE __NA__ category', async () => { + // The 105K dataset's `gene_name` carries `na` and `nan` as separate dictionary + // entries; they must land in a single legend slot, not two. + const payloads = { + ...PAYLOADS, + 'dict:organism': utf8('HumannaMousenan'), + 'dict:organism:len': i32(5, 2, 5, 3), + }; + const { data } = await decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) })); + + expect(data.annotations.organism.values).toEqual(['Human', 'Mouse', NA_VALUE]); + expect(data.annotations.organism.colors).toHaveLength(3); + // One NA slot, in the NA swatch — not a category sitting at the palette rank the + // token happened to have, and not a second slot beside the synthetic one. + expect(data.annotations.organism.colors.at(-1)).toBe(NA_DEFAULT_COLOR); + // Rows 1 and 5 spelled `na`, row 7 spelled `nan`, row 3 was already missing: all + // four end up on code 2. + expect(Array.from(data.annotation_data.organism as Int32Array)).toEqual([ + 0, 2, 1, 2, 0, 2, 1, 2, + ]); + }); + + it('drops a folded CSR hit with its score run and its evidence code', async () => { + // `binding` / `none` / `transport`: `none` is 1383 of 1587 rows of + // `phosphatase.predicted_transmembrane`, so this is the shipped shape. + const payloads = { + ...PAYLOADS, + 'dict:go_bp': utf8('bindingnonetransport'), + 'dict:go_bp:len': i32(7, 4, 9), + // Hit 1 is P2's `none` and now owns a score of 7, which must vanish with it + // while every later hit keeps its own. + 'score_count:go_bp': i32(2, 1, 1, 1, 0, 0, 0, 3, 0), + 'scores:go_bp': f64(1.5, 2.5, 7, 9.75, 4, 0.5, 0.25, 0.125), + }; + const { data } = await decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) })); + + expect(data.annotations.go_bp.values).toEqual(['binding', 'transport', NA_VALUE]); + expect(Array.from((data.annotation_data.go_bp as CsrAnnotationData).end)).toEqual([ + 1, 2, 3, 4, 6, 7, 9, 10, + ]); + expect(labelsOf(data, 'go_bp', 1)).toEqual(['binding']); + expect(labelsOf(data, 'go_bp', 4)).toEqual(['binding', 'transport']); + // P6's only hit was `none`, so the row empties out and takes the same __NA__. + expect(labelsOf(data, 'go_bp', 5)).toEqual([NA_VALUE]); + expect(getProteinScores(data, 1, 'go_bp')).toEqual([[1.5, 2.5]]); + expect(getProteinScores(data, 2, 'go_bp')).toEqual([[9.75]]); + expect(getProteinScores(data, 4, 'go_bp')).toEqual([[4], null]); + expect(getProteinScores(data, 6, 'go_bp')).toEqual([[0.5, 0.25, 0.125], null]); + expect(getProteinEvidence(data, 2, 'go_bp')).toEqual(['ECO:0000269']); + // P6's IDA went with its hit. + expect(getProteinEvidence(data, 5, 'go_bp')).toEqual([null]); + }); + it('leaves a multi column with neither scores nor evidence without those payloads', async () => { const { data } = await decodeParquetBundle(v3Bundle()); @@ -426,7 +496,7 @@ describe('parquetbundle format v3', () => { }); it('score counts that do not sum to the score count', async () => { - const payloads = { ...PAYLOADS, 'scores:go_bp': f32(1.5, 2.5) }; + const payloads = { ...PAYLOADS, 'scores:go_bp': f64(1.5, 2.5) }; await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( /score counts sum to 7 but scores:go_bp holds 2/, ); @@ -446,6 +516,15 @@ describe('parquetbundle format v3', () => { ); }); + it('a score payload whose byte length is not a multiple of 8', async () => { + // 12 bytes clears the int32 alignment the other payloads use, so only a + // float64-aware check catches it. + const payloads = { ...PAYLOADS, 'scores:go_bp': f32(1.5, 2.5, 9.75) }; + await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( + /"scores:go_bp" is 12 bytes, not a multiple of 8/, + ); + }); + it('an evidence code outside the evidence dictionary', async () => { const payloads = { ...PAYLOADS, 'evidence:go_bp': i32(-1, 0, 1, -1, -1, -1, 5, -1, -1) }; await expect(decodeParquetBundle(v3Bundle({ 5: payloadPart(payloads) }))).rejects.toThrow( diff --git a/packages/core/src/components/data-loader/utils/bundle-v3.ts b/packages/core/src/components/data-loader/utils/bundle-v3.ts index e423e7cb..1018cbc5 100644 --- a/packages/core/src/components/data-loader/utils/bundle-v3.ts +++ b/packages/core/src/components/data-loader/utils/bundle-v3.ts @@ -11,12 +11,17 @@ * This reader therefore never parses a string that is not a label, and hands the worker * typed arrays it can transfer instead of structured-clone. * - * Two wire details drive most of the code here: + * Three wire details drive most of the code here: * * - **Lengths are per-element counts, never cumulative offsets.** Offsets are * near-incompressible; their first differences are not. Every `__count`, * `score_count:` and `dict::len` family is prefix-summed here into the * cumulative offsets the in-memory `CsrAnnotationData` / `CsrScores` types use. + * - **The dictionaries are faithful, not presentational.** The encoder stopped + * collapsing `none`/`NA`/`null` because doing so corrupted the Python side, so + * `dict:` can carry those spellings as ordinary labels and this reader folds + * them into `__NA__` (see {@link foldMissingLabels}), exactly where the v2 path + * has always applied that rule. * - **Every part 1/3/6 column is REQUIRED and PLAIN**, which is the only shape * hyparquet decodes straight into a typed array. A column that arrives as a plain * array still reads correctly (see the fallback in {@link writeChunk}) but about 4x @@ -27,6 +32,7 @@ import { parquetMetadata, parquetRead, parquetReadObjects, type FileMetaData } f import { NA_DEFAULT_COLOR, NA_VALUE, + normalizeMissingValue, type Annotation, type AnnotationData, type BundleSettings, @@ -373,15 +379,18 @@ async function readPayloads(part: ArrayBuffer): Promise> * byte offset — so it is copied rather than wrapped. Wrapping would both risk an * alignment error and pin (or, if transferred, detach) the whole decoded page. */ -function asTypedPayload( +function asTypedPayload( payloads: ReadonlyMap, name: string, - ctor: new (buffer: ArrayBuffer) => T, + ctor: { new (buffer: ArrayBuffer): T; readonly BYTES_PER_ELEMENT: number }, ): T { const bytes = payloads.get(name); if (!bytes) throw new Error(`v3 bundle is missing the "${name}" payload`); - if (bytes.byteLength % 4 !== 0) { - throw new Error(`v3 payload "${name}" is ${bytes.byteLength} bytes, not a multiple of 4`); + const width = ctor.BYTES_PER_ELEMENT; + if (bytes.byteLength % width !== 0) { + throw new Error( + `v3 payload "${name}" is ${bytes.byteLength} bytes, not a multiple of ${width}`, + ); } return new ctor(bytes.slice().buffer); } @@ -421,6 +430,38 @@ function readLabels(payloads: ReadonlyMap, name: string): st return labels; } +/** + * Drop the dictionary entries that spell a missing value, exactly as v2 ingestion does. + * + * The encoder stopped collapsing `none`/`NA`/`null` (it corrupted the Python side — + * 1383 rows of `phosphatase.predicted_transmembrane` are literally the word `none`), so + * v3 is a faithful container and the presentation rule belongs here, which is where the + * v2 path has always applied it: `splitCategoricalAnnotationValues` filters these + * spellings out of a cell before anything counts frequencies, so a row left with nothing + * falls through to the same synthetic `__NA__` v2 gives it. That is also why folding + * cannot produce a second NA slot: `__na__` is itself a missing-value token, so + * `NA_VALUE` never survives this pass and the append below is the only NA there is. + * + * Compaction preserves the survivors' relative order, so the encoder's + * descending-frequency dictionary order — and with it the palette assignment — is + * unchanged; the colours are generated from the post-fold length. + * + * Returns the old-code -> new-code map (`-1` for a dropped entry), or `null` when the + * dictionary is already clean. `labels` is compacted in place. + */ +function foldMissingLabels(labels: string[]): Int32Array | null { + const remap = new Int32Array(labels.length); + let kept = 0; + for (let i = 0; i < labels.length; i++) { + const drop = normalizeMissingValue(labels[i]) === null; + remap[i] = drop ? -1 : kept; + if (!drop) labels[kept++] = labels[i]; + } + if (kept === labels.length) return null; + labels.length = kept; + return remap; +} + /** Per-element counts to the cumulative offsets the in-memory CSR types use. */ function prefixSum(counts: Int32Array, what: string): Int32Array { const end = new Int32Array(counts.length); @@ -434,10 +475,22 @@ function prefixSum(counts: Int32Array, what: string): Int32Array { return end; } +/** + * Score values are read as **float64**, matching what the encoder writes. + * + * float32 destroys the E-value, which is the canonical Pfam and InterPro score: 1e-200 + * flushes to 0 and 1e40 saturates to Infinity. `CsrScores.values` in `@protspace/utils` + * is still declared `Float32Array` and has to be widened to accept this. + */ +interface V3Scores { + hitEnd: Int32Array; + values: Float64Array; +} + interface CsrColumn { end: Int32Array; codes: Int32Array; - scores: CsrScores | null; + scores: V3Scores | null; evidence: CsrEvidence | null; } @@ -466,7 +519,7 @@ function readCsrColumn( } } - let scores: CsrScores | null = null; + let scores: V3Scores | null = null; if (column.scores) { const scoreCounts = asTypedPayload(payloads, `score_count:${name}`, Int32Array); if (scoreCounts.length !== codes.length) { @@ -474,7 +527,7 @@ function readCsrColumn( `v3 column "${name}" has ${scoreCounts.length} score counts for ${codes.length} hits`, ); } - const values = asTypedPayload(payloads, `scores:${name}`, Float32Array); + const values = asTypedPayload(payloads, `scores:${name}`, Float64Array); const hitEnd = prefixSum(scoreCounts, `column "${name}" score counts`); const scoreTotal = hitEnd.length > 0 ? hitEnd[hitEnd.length - 1] : 0; if (scoreTotal !== values.length) { @@ -509,6 +562,59 @@ function readCsrColumn( return { end, codes, scores, evidence }; } +/** + * Renumber a CSR column onto a folded dictionary, dropping the hits whose label went + * with it — along with that hit's score run and evidence code, both of which are + * numbered by hit. A row left with nothing is picked up by {@link insertNAForEmptyRows} + * below, which is what v2 does with a cell whose only value was a missing-value + * spelling. + */ +function dropFoldedHits(csr: CsrColumn, remap: Int32Array | null): CsrColumn { + if (!remap) return csr; + + const numRows = csr.end.length; + const scores = csr.scores; + let keptHits = 0; + let keptScores = 0; + for (let hit = 0; hit < csr.codes.length; hit++) { + if (remap[csr.codes[hit]] < 0) continue; + keptHits++; + if (scores) keptScores += scores.hitEnd[hit] - (hit === 0 ? 0 : scores.hitEnd[hit - 1]); + } + + const codes = new Int32Array(keptHits); + const end = new Int32Array(numRows); + const evidenceCodes = csr.evidence ? new Int32Array(keptHits) : null; + const hitEnd = scores ? new Int32Array(keptHits) : null; + const values = scores ? new Float64Array(keptScores) : null; + + let write = 0; + let scoreWrite = 0; + for (let row = 0; row < numRows; row++) { + for (let hit = row === 0 ? 0 : csr.end[row - 1]; hit < csr.end[row]; hit++) { + const code = remap[csr.codes[hit]]; + if (code < 0) continue; + codes[write] = code; + if (evidenceCodes) evidenceCodes[write] = csr.evidence!.codes[hit]; + if (hitEnd) { + for (let at = hit === 0 ? 0 : scores!.hitEnd[hit - 1]; at < scores!.hitEnd[hit]; at++) { + values![scoreWrite++] = scores!.values[at]; + } + hitEnd[write] = scoreWrite; + } + write++; + } + end[row] = write; + } + + return { + end, + codes, + scores: scores ? { hitEnd: hitEnd!, values: values! } : null, + evidence: csr.evidence ? { codes: evidenceCodes!, dict: csr.evidence.dict } : null, + }; +} + /** * Route rows with no hits at all to a synthetic `__NA__` category, the way * `appendSyntheticNACategory` does for the nested storage shape. @@ -638,28 +744,36 @@ export async function readV3Bundle( } const labels = readLabels(payloads, name); + // Codes on the wire index the dictionary AS WRITTEN, so they are range-checked + // against that count and only then renumbered onto the folded one. + const encodedLabelCount = labels.length; + const remap = foldMissingLabels(labels); const { colors, shapes } = generateColorsAndShapes('kellys', labels.length); if (column.kind === 'categorical') { const codes = stored as Int32Array; for (let i = 0; i < numRows; i++) { - if (codes[i] >= labels.length || codes[i] < -1) { + if (codes[i] >= encodedLabelCount || codes[i] < -1) { throw new Error( - `v3 column "${name}" row ${i} has code ${codes[i]}, outside its ${labels.length} labels`, + `v3 column "${name}" row ${i} has code ${codes[i]}, outside its ${encodedLabelCount} labels`, ); } + if (remap && codes[i] >= 0) codes[i] = remap[codes[i]]; } appendSyntheticNACategoryToCodes(labels, colors, shapes, codes); annotation_data[name] = codes; } else { const csr = insertNAForEmptyRows( - readCsrColumn( - name, - column, - stored as Int32Array, - labels.length, - payloads, - readEvidenceDict, + dropFoldedHits( + readCsrColumn( + name, + column, + stored as Int32Array, + encodedLabelCount, + payloads, + readEvidenceDict, + ), + remap, ), labels, colors, diff --git a/packages/utils/src/types.ts b/packages/utils/src/types.ts index dc111d45..4cf76ae8 100644 --- a/packages/utils/src/types.ts +++ b/packages/utils/src/types.ts @@ -76,7 +76,12 @@ export interface CsrAnnotationData { */ export interface CsrScores { readonly hitEnd: Int32Array; - readonly values: Float32Array; + /** + * float64, matching the `scores:` payload the v3 encoder writes. float32 + * cannot carry an E-value — the canonical Pfam / InterPro score — at all: 1e-200 + * flushes to 0 and 1e40 saturates to Infinity. + */ + readonly values: Float64Array; } /** diff --git a/packages/utils/src/visualization/eat-overlay.test.ts b/packages/utils/src/visualization/eat-overlay.test.ts index 289728aa..68928abd 100644 --- a/packages/utils/src/visualization/eat-overlay.test.ts +++ b/packages/utils/src/visualization/eat-overlay.test.ts @@ -262,7 +262,7 @@ describe('EAT overlay over CSR storage (bundle format v3)', () => { annotation_scores_csr: { ec: { hitEnd: Int32Array.of(1, 2, 3, 4, 4), - values: Float32Array.of(0.25, 0.5, 0.75, 0.125), + values: Float64Array.of(0.25, 0.5, 0.75, 0.125), }, }, annotation_evidence_csr: { @@ -302,7 +302,9 @@ describe('EAT overlay over CSR storage (bundle format v3)', () => { expect(Array.from(data.annotation_evidence_csr!.ec.codes)).toEqual([0, 1, 2, 0, -1]); // p1's 0.5 is gone with the row it belonged to; the slack is not retained. expect(Array.from(out.annotation_scores_csr!.ec.values)).toEqual([0.25, 0.75, 0.125]); - expect(out.annotation_scores_csr!.ec.values.buffer.byteLength).toBe(3 * 4); + expect(out.annotation_scores_csr!.ec.values.buffer.byteLength).toBe( + 3 * Float64Array.BYTES_PER_ELEMENT, + ); }); it('does not grow the payload records on a dataset that carries none', () => { diff --git a/packages/utils/src/visualization/eat-overlay.ts b/packages/utils/src/visualization/eat-overlay.ts index f5664493..fe8c7310 100644 --- a/packages/utils/src/visualization/eat-overlay.ts +++ b/packages/utils/src/visualization/eat-overlay.ts @@ -255,7 +255,7 @@ function cloneCsrWithPredictions( const codes = new Int32Array(total); // The payloads can only shrink (a replaced row drops its own values), so the source // length is a safe upper bound and the trailing slack is sliced off at the end. - const values = new Float32Array(sourceScores ? sourceScores.values.length : 0); + const values = new Float64Array(sourceScores ? sourceScores.values.length : 0); const hitEnd = new Int32Array(sourceScores ? total : 0); const evidenceCodes = new Int32Array(sourceEvidence ? total : 0); diff --git a/packages/utils/src/visualization/plot-data-accessors.test.ts b/packages/utils/src/visualization/plot-data-accessors.test.ts index ba6b4649..0a2f0fbe 100644 --- a/packages/utils/src/visualization/plot-data-accessors.test.ts +++ b/packages/utils/src/visualization/plot-data-accessors.test.ts @@ -427,7 +427,7 @@ describe('CSR score and evidence payloads', () => { // hit 0 -> [1.5]; hit 1 -> no scores; hit 2 -> [0.25, 0.5] const csrScores: CsrScores = { hitEnd: Int32Array.from([1, 1, 3]), - values: Float32Array.from([1.5, 0.25, 0.5]), + values: Float64Array.from([1.5, 0.25, 0.5]), }; const csrEvidence: CsrEvidence = { codes: Int32Array.from([0, -1, 1]), diff --git a/packages/utils/src/visualization/slice-visualization-data.test.ts b/packages/utils/src/visualization/slice-visualization-data.test.ts index c804aba8..46b207d3 100644 --- a/packages/utils/src/visualization/slice-visualization-data.test.ts +++ b/packages/utils/src/visualization/slice-visualization-data.test.ts @@ -114,7 +114,7 @@ describe('sliceVisualizationDataByIndices over CSR storage (bundle format v3)', annotation_scores_csr: { fam: { hitEnd: Int32Array.of(1, 1, 3, 4), - values: Float32Array.of(0.5, 1.5, 2.5, 3.5), + values: Float64Array.of(0.5, 1.5, 2.5, 3.5), }, }, annotation_evidence_csr: { diff --git a/packages/utils/src/visualization/slice-visualization-data.ts b/packages/utils/src/visualization/slice-visualization-data.ts index a30a4f6a..bbf4f51d 100644 --- a/packages/utils/src/visualization/slice-visualization-data.ts +++ b/packages/utils/src/visualization/slice-visualization-data.ts @@ -64,7 +64,7 @@ export function sliceVisualizationDataByIndices( let total = 0; for (const hit of hits) total += csr.hitEnd[hit] - (hit === 0 ? 0 : csr.hitEnd[hit - 1]); const hitEnd = new Int32Array(hits.length); - const values = new Float32Array(total); + const values = new Float64Array(total); let cursor = 0; for (let k = 0; k < hits.length; k++) { const hit = hits[k]; From 6db1ec70e0eaecbc2e9e81ece9cee645306e2563 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 02:55:21 +0200 Subject: [PATCH 25/31] docs(bundle): document the six-part v3 container and its physical schema Retitle the v2 cell grammar as the logical layer and add version detection for 1, 2 and 3. --- apps/protspace/CLAUDE.md | 7 +- docs/guide/data-format.md | 267 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 260 insertions(+), 14 deletions(-) diff --git a/apps/protspace/CLAUDE.md b/apps/protspace/CLAUDE.md index d9c3fbec..efc85195 100644 --- a/apps/protspace/CLAUDE.md +++ b/apps/protspace/CLAUDE.md @@ -248,8 +248,13 @@ HDF5 file (float16 embeddings) 3. `projections_data` — reduced coordinates per protein per projection 4. `settings` (optional) — annotation styles, pinned values, display config 5. `statistics` (optional) — tidy table of annotation-based validity (silhouette/DBI/CH per annotation, `space_kind ∈ {embedding, projection}`, `annotation` column) + auto-cluster ARI/NMI agreement (`stat_family=cluster_agreement`) (`protspace stats` / `prepare --stats`) +6. `payloads` (format v3 only, required) — label dictionaries and CSR code/score/evidence buffers for part 1 (`data/io/bundle_v3.py`) -Positional layout is `core(3) + settings? + statistics?`. When statistics are present but settings are absent, the settings slot is written as **zero bytes** so statistics stay at position five (readers branch on emptiness, not part count). Both bundled and separate-file (`--no-bundled`) output persist `settings.parquet` and `statistics.parquet` when present. +Every write from here emits **six** parts (format v3): `core(3) + settings + statistics + payloads`, with zero bytes in the settings or statistics slot when absent, because the browser reads the payloads positionally from `parts[5]`. `replace_settings_in_bundle` (`protspace style`) is the exception: it preserves the layout it was given, so a legacy bundle stays legacy. + +Legacy (v1/v2) containers still read: positional layout `core(3) + settings? + statistics?`, 3 to 5 parts. When statistics are present but settings are absent, the settings slot is written as **zero bytes** so statistics stay at position five (readers branch on emptiness, not part count). Both bundled and separate-file (`--no-bundled`) output persist `settings.parquet` and `statistics.parquet` when present. + +`read_tables()` / `read_bundle()` / `extract_bundle_to_dir()` decode v3 back to the v2-shaped tables (all-string cells, long projections, footer re-stamped `protspace_format_version=2`), so every consumer above `data/io/` is unchanged. `BUNDLE_FORMAT_VERSION = 2` versions the cell grammar, not the container. See `docs/guide/data-format.md`. ## Testing diff --git a/docs/guide/data-format.md b/docs/guide/data-format.md index b5b86fc6..1a36be5f 100644 --- a/docs/guide/data-format.md +++ b/docs/guide/data-format.md @@ -4,7 +4,24 @@ ProtSpace uses `.parquetbundle` files - a single file containing all visualizati ## What is a .parquetbundle? -A `.parquetbundle` is a single file containing three core Parquet tables bundled together, with optional settings and statistics sections: +A `.parquetbundle` is a single file that concatenates several Parquet files, separated by the +byte string `---PARQUET_DELIMITER---`. It keeps everything in one convenient file while still +loading efficiently in the browser. + +There are two container layouts. Which one a file uses is recorded in the Parquet key-value +metadata of its first part, under `protspace_format_version` (see +[Version detection](#version-detection)): + +| Layout | Parts | Written by | +| ------------------------- | -------- | -------------------------------------------------------------------------------------------------------------------- | +| Legacy (format v1 and v2) | 3 to 5 | every export from the web app; `protspace style`, which keeps the layout of the bundle it was handed; older releases | +| Columnar (format v3) | always 6 | `protspace prepare`, `protspace bundle`, `protspace transfer` | + +Both layouts carry the same data. v3 re-encodes the container, not the dataset: the Python API +decodes a v3 file back into exactly the three tables, with exactly the cell grammar, that a +legacy file stores directly. See [Format v3 Physical Schema](#format-v3-physical-schema). + +### Legacy layout (3 to 5 parts) ``` .parquetbundle file @@ -23,20 +40,47 @@ Part positions are fixed, not counted. Once a statistics part is present, the se and the settings slot are mandatory, even when there are no settings to store: the slot is then written as zero bytes so statistics stay at position five. +### Format v3 layout (always 6 parts) + +``` +.parquetbundle file (format v3) +├── selected_annotations.parquet # Integer codes, per-row hit counts, float64 numerics +├── ---PARQUET_DELIMITER--- # Separator +├── projections_metadata.parquet # Unchanged from the legacy layout +├── ---PARQUET_DELIMITER--- # Separator +├── projections_data.parquet # Wide float32 columns, one row per protein +├── ---PARQUET_DELIMITER--- # Separator +├── settings.parquet # Zero bytes when there are no settings +├── ---PARQUET_DELIMITER--- # Separator +├── statistics.parquet # Zero bytes when there are no statistics +├── ---PARQUET_DELIMITER--- # Separator +└── payloads.parquet # Required: label dictionaries and CSR buffers +``` + +All six slots are always emitted, and parts 4 and 5 are written as zero bytes when the bundle +carries no settings or no statistics. Part 6 is required and never empty. The slots are +positional rather than counted because the browser reads the payloads from a fixed index: an +omitted settings or statistics slot would file the payloads where statistics are expected, and +the reader would report a bundle with no payloads part. + This bundled format allows efficient loading in the browser while keeping everything in one convenient file. The optional settings section is stored as `settings.parquet`, a one-row Parquet table with a `settings_json` column. It stores legend customizations (colors, shapes, ordering, visibility, palette, numeric binning settings) and export options (image dimensions, legend sizing) per annotation. When present, these settings are applied automatically on load so the visualization renders exactly as it was exported. ## Tables +These are the logical tables every bundle carries. A legacy container stores them exactly as +described here. A v3 container stores an equivalent columnar encoding and decodes back to these +same tables on every Python read, so the shapes below are what a Python consumer always sees. + ### 1. Annotations Table Contains metadata and biological annotations for each protein. -| Column | Type | Description | -| ------------ | ------------- | -------------------------- | -| `identifier` | string | Protein ID (e.g., P12345) | -| _others_ | string/number | Any biological annotations | +| Column | Type | Description | +| ------------ | ------------- | ------------------------------------------------------------------ | +| `identifier` | string | Protein ID (e.g., P12345); named `protein_id` in Python CLI output | +| _others_ | string/number | Any biological annotations | The columns `gene_name`, `protein_name`, and `uniprot_kb_id` are **tooltip-only**, shown on hover but excluded from the annotation dropdown. @@ -134,6 +178,16 @@ become the N/A legend entry. There is no density, cardinality, or sparsity heuristic. +In a v3 bundle the scan runs once, at write time, and the answer is recorded in the manifest, so +the browser never re-scans the column. The write-time test differs from the browser's in two +small ways: + +- a column that is already numeric in Arrow stays numeric whatever it holds, including a column + with no values left at all, where the browser would fall back to categorical +- only decimal literals count as numbers. JavaScript also reads `0x10`, `0o17` and `0b1` as + numbers, so a column of such literals is categorical in a v3 file and numeric on the legacy + path. None of the shipped datasets contain one. + ::: warning Identifier-style number columns become numeric A column of cluster IDs or numeric codes stored as strings (`"1"`, `"2"`, `"17"`) parses cleanly as numbers, so it **is** treated as numeric and binned with a gradient, even if it is sparse and you @@ -182,6 +236,11 @@ How missing values are stored: on export a missing categorical cell is written a not as a sentinel string. Bundles from older web builds may still hold literal `__NA__` cells; those normalize to N/A on load, as above. +A v3 bundle stores a missing cell as code `-1`, a hit count of `0`, or `NaN`, depending on the +column. It never rewrites the spellings above into a missing value: they stay in the file as +ordinary labels and the reader folds them into N/A, exactly as it does for a legacy bundle. See +[Missing values in a v3 file](#missing-values-in-a-v3-file). + The single "N/A" legend row covers every missing-value protein. Its default color is light grey (`#DDDDDD`) and circle shape, matching every other category in the system. For categorical annotations the color and shape are @@ -212,7 +271,13 @@ Evidence codes are recognized by pattern: any 2–5 uppercase letter code (e.g., Evidence codes are displayed in the protein tooltip alongside the annotation value. -### Encoding (Format v2) +### Cell Grammar and Encoding (Format v2) {#encoding-format-v2} + +This section describes the **logical** cell grammar: the string spelling of an annotation value, +with its `;` separated values and its `|` suffixed scores and evidence codes. It is what a legacy +bundle stores on disk, and it is also what every Python consumer sees, because a v3 bundle +decodes back into exactly this grammar on every read. Only the physical storage differs, and only +inside a v3 container. As of bundle format v2, annotation values containing special characters use percent-encoding to ensure reliable parsing while keeping `,` `(` `)` human-readable inside names and labels. @@ -235,13 +300,6 @@ For a hypothetical protein with a CATH domain whose name contains a semicolon (" When displayed in ProtSpace, the decoded names render as "Superfamily; old" and "Kinase, serine", with the percent-encoding transparent to the user. -**Version detection:** - -- A bundle's annotation format version is stored in the parquet key-value metadata of the `selected_annotations` table under the key `protspace_format_version` -- Format version 2 is detected by reading this metadata via hyparquet's `parquetMetadata` (returns `"2"` as a string) -- v1 bundles (no version key present, or version < 2) render using the legacy parser, which does not decode percent-encoded sequences -- This ensures backward compatibility: existing v1 bundles load unchanged without requiring special-case handling - **Numeric column typing:** - A numeric Parquet type does **not** by itself make a column numeric. Numeric-ness is decided by the content scan described in [Numeric Annotations](#numeric-annotations); the declared type is consulted in only two cases, to upgrade an already-numeric column's type to `int`, and to rescue a column whose every row is missing. A column with any real value stays categorical whatever the schema says. @@ -249,11 +307,194 @@ When displayed in ProtSpace, the decoded names render as "Superfamily; old" and - Only an _unannotated_ physical type counts. A logical or converted type means the physical type is a carrier rather than the identity (pyarrow stores an all-null column as `INT32` + logical `NULL`, and `DECIMAL` rides on `INT32`/`INT64`), so those fall back to inference too - This matters for a column whose rows are all missing, for example an isolation-mode or query-filtered export. Such a column has no values left to infer from and would otherwise reload as a categorical column with a single N/A category - Bundles that store annotations as text (the `protspace` CLI writes its annotation frame stringified) carry no such type, and fall back to inferring numeric-ness from the values as before +- All of this applies to legacy bundles only. In a v3 bundle the manifest declares each column's kind and `numericType`, and the physical Parquet types are deliberately not consulted, because every dictionary-code column there is an `INT32` **Known formatting:** - Unnamed CATH superfamilies from TED domains display the bare code without a decoding step (see [tsenoner/protspace-legacy#57](https://github.com/tsenoner/protspace-legacy/issues/57)) +## Version Detection + +- A bundle's format version lives in the Parquet key-value metadata of the `selected_annotations` + part, under the key `protspace_format_version`, and is read from that part's footer before any + row is decoded. +- `"2"` selects the percent-decoding cell parser described above, `"3"` selects the columnar + reader. A v1 bundle has no version key and renders with the legacy parser, which does not + decode percent-encoded sequences. Existing v1 and v2 bundles therefore keep loading unchanged. +- Python cross-checks the two signals: a six-part file whose first part does not say `3` is + rejected rather than guessed at, and a three to five part file is always read as legacy. +- The key versions the **container**, not the cell grammar. The grammar is still v2, which is why + `BUNDLE_FORMAT_VERSION` on the Python side stays `2` and why the tables handed back from a v3 + read are re-stamped `protspace_format_version=2`: what they contain is v2 cells. +- A web build older than v3 support rejects a v3 file with + `Expected 2 to 4 delimiters in parquetbundle, found 5`. That is a version-skew signal, not a + corrupt file. + +## Format v3 Physical Schema + +Format v3 changes how the three logical tables are physically stored so that the browser can +build its typed arrays without parsing a single annotation cell. On the 573K SwissProt dataset +that takes bundle decoding from about 6.5 s and 2.1 GB of heap to about 0.4 s and under 50 MB, +in a file about 19% smaller than the v2 encoding of the same data. + +Nothing above the container boundary changes. `read_tables()` decodes a v3 file back into the +v2-shaped tables, so `protspace serve`, `protspace style`, `protspace transfer` and every script +keep their string-cell logic. + +### Part 1: annotations + +One row per protein, the identifier column first and then the annotation columns in their +original order. The manifest says how to read each one: + +| Manifest kind | Part 1 column | Physical type | Holds | Missing | +| ------------- | ---------------------------- | ------------- | -------------------------------------------------- | -------------- | +| (identifier) | `protein_id` or `identifier` | `BYTE_ARRAY` | the protein ID; null and duplicate IDs are refused | not allowed | +| `categorical` | `` | `INT32` | a code into `dict:` | `-1` | +| `multi` | `__count` | `INT32` | how many hits this row owns in `csr:` | a count of `0` | +| `numeric` | `` | `DOUBLE` | the value | `NaN` | + +A column is `multi` when any row holds more than one value, or when any value carries a score or +an evidence code. Part 1 then stores only the per-row hit counts; the hits themselves live in +part 6. The EAT companion columns (`__pred_value`, `__pred_confidence`, +`__pred_source`) follow the same rules by kind. + +Part 2, the projections metadata, is unchanged from the legacy layout. + +### Part 3: projections + +One row per protein, aligned position by position with part 1, and one `FLOAT` column per axis: +`__x`, `__y`, and `__z` for a 3D projection. A protein with no coordinates in a +projection is stored at `0.0`, which is where the legacy reader placed it too. + +The projection name sets in parts 2 and 3 must be identical; the encoder refuses a bundle where +either one names a projection the other does not, because the browser derives the projection set +from the data rows alone. + +### Part 6: payloads + +A two-column table, `name` (string) and `data` (binary), one row per payload. Every payload is a +raw little-endian buffer: + +| Payload | Element | One per | Contents | +| ---------------------------------------- | ---------- | ------- | ------------------------------------------------------------- | +| `dict:` | utf8 bytes | blob | every label of `` concatenated, in code order | +| `dict::len` | int32 | label | that label's length in bytes | +| `csr:` | int32 | hit | the label code of each hit, rows in row order | +| `score_count:` | int32 | hit | how many score values that hit owns | +| `scores:` | float64 | score | the score values, in hit order | +| `evidence:` | int32 | hit | index into `dict:__evidence`, `-1` for a hit with no evidence | +| `dict:__evidence`, `dict:__evidence:len` | as above | | the single evidence dictionary every column indexes into | + +Row `i` of a `multi` column owns the codes at `[start, start + count)`, where `start` is the sum +of the counts of every row before it. + +Labels are stored **decoded**: the percent-encoding is removed before the dictionary is written, +and applied again when Python decodes the bundle. Codes run in descending order of how many hits +carry the label, ties broken by first occurrence, so code `0` is the column's most frequent label +and the palette lands on the same categories it did in v2. + +### Counts, never offsets + +Every length family in v3 is a per-element **count**, never a cumulative offset: `__count` +counts a row's hits, `score_count:` counts a hit's score values, `dict::len` counts a +label's bytes. The reader prefix-sums them into offsets in a single pass. + +That is a size decision. Offsets are near incompressible (snappy manages about 0.4% of them on +the 573K SwissProt bundle) while their first differences, which is what the counts are, compress +about 8x. On that bundle the difference is roughly 15 MB, about 10.0 MB of it in part 1 and +about 5.7 MB in part 6. + +### Required, PLAIN, one row group + +Every column of parts 1, 3 and 6 is written non-nullable, PLAIN encoded, with dictionary encoding +disabled, in a single row group, snappy compressed. + +That is load bearing, not stylistic. The browser's Parquet reader hands back a zero-copy typed +array only for a REQUIRED flat PLAIN column. A column written nullable or dictionary-encoded +still decodes to the right values, but it arrives as a plain JavaScript array about 4x slower, +and the reader logs one warning naming it. Such a column is a writer bug, not a variant. + +### The manifest + +Part 1's footer carries the two key-value entries that describe the format: +`protspace_format_version`, which is `"3"`, and `protspace_v3_manifest`, a JSON object that is +the only description of what the integer columns mean. + +```json +{ + "idColumn": "protein_id", + "columns": { + "family": { "kind": "multi", "scores": true, "sourceType": "string" }, + "length": { "kind": "numeric", "numericType": "int", "sourceType": "int32" } + }, + "projections": [{ "name": "UMAP_2", "dimension": 2 }] +} +``` + +| Field | Meaning | +| ----------------------- | ------------------------------------------------------------------------------ | +| `idColumn` | which part 1 column holds the protein IDs | +| `columns.*.kind` | `categorical`, `multi` or `numeric` | +| `columns.*.numericType` | `int` or `float`; numeric columns only | +| `columns.*.scores` | present and `true` when a `scores:` payload exists; multi columns only | +| `columns.*.evidence` | present and `true` when an `evidence:` payload exists; multi columns only | +| `columns.*.sourceType` | Python-private, see below | +| `projections` | `{ name, dimension }` per projection, in part 3 column order | + +The browser validates the manifest against part 1's own schema before reading anything. An +unknown kind, a declared column part 1 does not have, a kind whose physical type disagrees (a +code column declared numeric, for instance), a duplicate projection name, or a dimension that is +not 2 or 3 all throw rather than being repaired. + +`sourceType` is Python-private and the browser ignores it. It records the Arrow type the column +had before encoding, so the decoder can restore it instead of handing back a string column, and +it is `"?"` for a type that cannot be parsed back from its alias, such as a dictionary, list or +decimal column. The decoder falls back to the per-kind default for those. + +### Scores are float64 + +`scores:` is the one wide payload in an otherwise narrow format. float32 cannot carry an +E-value, which is the canonical Pfam and InterPro score: `1e-200` flushes to zero and `1e40` +overflows to infinity, and infinity is not a valid v2 cell, so a second round trip would +reclassify the hit as a plain label. On the 573K SwissProt bundle the float64 scores cost about +940 KB. + +### Missing values in a v3 file {#missing-values-in-a-v3-file} + +Only a null cell and an empty cell (after trimming whitespace) are missing in v3. The spellings +listed under [Missing Values](#missing-values), `none`, `NA`, `n/a`, `nan`, `null` and `__NA__`, +are kept in the file as ordinary labels, and the browser folds them into its N/A category at read +time, on v3 exactly as it always has on v2. + +The reason is that v3 is a container encoding and must hand back the label it was given. +Collapsing these spellings in the file broke `protspace style` on the shipped phosphatase +dataset, where 1383 of 1587 rows of `predicted_transmembrane` are literally the word `none`: the +style command raised, and a 1383-protein legend entry came back blank. Folding them stays a +display decision, made by the reader. + +Those spellings are consulted at write time in one place only, to decide whether a column is +numeric, so a column of `NA` stays categorical instead of becoming an all-`NaN` numeric column. + +### What a v3 round trip does not preserve + +A v3 file stores what the reader would have parsed out of the v2 cells, not the cells themselves, +so decoding a v3 bundle returns the canonical spelling of each cell rather than the original +bytes. The differences are deliberate: + +| Written | Read back | Why | +| ------------------------------------- | --------------- | ------------------------------------------------------------------------------------------------- | +| `PF00001\|62.0` | `PF00001\|62` | scores are re-spelled the way JavaScript prints them, and JavaScript never prints a trailing `.0` | +| `PF00001\|0.5700` | `PF00001\|0.57` | same rule: the shortest spelling that reads back as the same double | +| ` A \|IDA` | `A\|IDA` | cells and hits are whitespace-trimmed | +| `A;;B` | `A;B` | empty hits are dropped | +| `%3b` | `%3B` | labels are re-encoded canonically, in uppercase hex | +| a missing cell | `""` | null and blank both mean missing | +| `100.0` where every value is integral | `100` | the canonical v2 spelling of an integral value | + +A cell spelled `none`, `NA` or `null` is an ordinary label and comes back unchanged. Projection +coordinates come back as float32 with `z` null for a 2D projection, and the identifier column +comes back first whatever position it held before. + ## Creating Files Use the [Google Colab notebook](/guide/data-preparation) or [Python CLI](/guide/python-cli) to generate `.parquetbundle` files. From b9c691fb3679fffb04b79a396e25c4dd0be65f80 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 03:01:07 +0200 Subject: [PATCH 26/31] test(bundle): add the golden v3 fixture both languages read Generated by scripts/generate_v3_fixture.py; a superset of v2-sample. --- apps/protspace/scripts/generate_v3_fixture.py | 188 +++++++++++++++++ .../protspace/tests/test_bundle_v3_fixture.py | 198 ++++++++++++++++++ .../__fixtures__/v3-sample.parquetbundle | Bin 0 -> 12809 bytes 3 files changed, 386 insertions(+) create mode 100644 apps/protspace/scripts/generate_v3_fixture.py create mode 100644 apps/protspace/tests/test_bundle_v3_fixture.py create mode 100644 packages/core/src/components/data-loader/utils/__fixtures__/v3-sample.parquetbundle diff --git a/apps/protspace/scripts/generate_v3_fixture.py b/apps/protspace/scripts/generate_v3_fixture.py new file mode 100644 index 00000000..c3885865 --- /dev/null +++ b/apps/protspace/scripts/generate_v3_fixture.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +"""Regenerate the browser's golden format-v3 bundle fixture. + +``packages/core/src/components/data-loader/utils/__fixtures__/v3-sample.parquetbundle`` +is the cross-language contract for parquetbundle v3: Python writes it here, +`bundle-v3.ts` reads it in vitest. It is a **superset** of the committed +``v2-sample.parquetbundle``: proteins ``P1``/``P2`` carry byte-identical +``cath`` and ``go_bp`` cells (including the percent-encoded ``;`` inside a CATH +name and the ``|IDA`` evidence suffix), so every assertion the v2 golden test +makes still holds, and four more proteins plus six more columns cover the v3 +paths a two-row two-column table cannot reach: + +* a plain single-valued categorical (``kingdom``) whose frequency order differs + from its first-occurrence order; +* multi-valued columns with scores (``cath``, ``pfam``) and with evidence codes + (``go_bp``, both the ``IDA`` and the ``ECO:0000269`` spellings); +* ``pfam`` also has two scores on one hit, and zero hits at the first, an + interior and the last row, so the reader's CSR prefix-sum and the synthetic + ```` insertion are exercised at every boundary; +* scores that only survive in float64: ``1e-200`` flushes to zero in float32 + and ``123456789`` re-spells as ``1.2345679e+08``; +* a numeric int column (``length``) and a numeric float one + (``hydrophobicity``), each with a blank cell; +* ``predicted_tm``, whose labels are the literal missing-value spellings + ``none`` and ``NA`` — Python keeps them (v3 is a container encoding), the + browser folds them into ```` at read time; +* a 2D (``pca2``, P1/P2 at the v2 fixture's coordinates) and a 3D (``umap3``) + projection; +* the EAT companion trio on ``kingdom`` (``__pred_value`` string, + ``__pred_confidence`` float32, ``__pred_source`` string), null for the + proteins with no prediction. + +Settings and statistics are deliberately absent, so the container's two +zero-byte slots keep the payloads part at position six. + +Usage:: + + cd apps/protspace && uv run python scripts/generate_v3_fixture.py +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa + +from protspace.data.annotations.encoding import encode_field, stamp_format_version +from protspace.data.io.bundle import write_bundle +from protspace.data.io.predictions import add_overlay_columns +from protspace.data.processors.base_processor import BaseProcessor + +FIXTURE_PATH = ( + Path(__file__).resolve().parents[3] + / "packages" + / "core" + / "src" + / "components" + / "data-loader" + / "utils" + / "__fixtures__" + / "v3-sample.parquetbundle" +) + +PROTEIN_IDS = ["P1", "P2", "P3", "P4", "P5", "P6"] + +# The v2 fixture's CATH name: a label whose own text contains the ';' the v2 +# grammar reserves as the hit separator, so it has to travel percent-encoded. +_CATH_NAME = encode_field("Ribosomal Protein L15; Chain: K; domain 2") +_CATH_1 = f"G3DSA:1.10.10.10 ({_CATH_NAME})" + +ANNOTATION_CELLS: dict[str, list[str]] = { + # P1/P2 are the v2 fixture's cells, verbatim. + "cath": [ + f"{_CATH_1}|50.2;G3DSA:6.20.10.10|60.5", + "6.20.10.10", + "G3DSA:6.20.10.10|123456789", + "", + f"{_CATH_1}|1e-200", + "6.20.10.10", + ], + "go_bp": [ + "apoptotic process|IDA", + "", + "apoptotic process|IDA;protein folding|ECO:0000269", + "protein folding|IEA", + "", + "apoptotic process|EXP", + ], + # Zero hits first, interior and last; two scores on one hit; a label + # carrying the encoded '|' the grammar reserves as the suffix separator. + "pfam": [ + "", + "PF00001 (7tm%3B1)|1e-10,2.5;PF00002|0.5", + "", + "PF00001 (7tm%3B1)|0.25", + "PF00003 (a%7Cb)|3", + "", + ], + # Frequency order (Bacteria, Archaea, Eukaryota) differs from first + # occurrence only at the tail, which is what pins the tie-break rule. + "kingdom": ["Bacteria", "Archaea", "Bacteria", "Eukaryota", "Bacteria", "Archaea"], + # Literal missing-value spellings kept as labels: a display decision the + # browser makes, not a container one. + "predicted_tm": ["none", "none", "TM helix", "none", "NA", "TM helix"], + "length": ["120", "", "340", "0", "-15", "1024"], + "hydrophobicity": ["0.5", "-1.25", "", "3.0", "1e-3", "42"], +} + +# (query_id, label, reliability, distance, source_id) for the EAT overlay. +PREDICTIONS = [ + ("P2", "Bacteria", 0.875, 0.12, "Q9XYZ1"), + ("P4", "Archaea", 0.5, 0.44, "P0A7B8"), + ("P5", "Bacteria", 0.25, 0.91, "A0A123"), +] + +PROJECTIONS = [ + { + "name": "pca2", + "dimensions": 2, + "info": {"components": 2}, + # P1 and P2 sit exactly where the v2 fixture puts them. + "data": np.array( + [ + [0.0, 0.0], + [1.0, 1.0], + [2.5, -3.5], + [-4.0, 0.25], + [5.0, 5.0], + [-1.5, 2.0], + ], + dtype=np.float32, + ), + }, + { + "name": "umap3", + "dimensions": 3, + "info": {"n_neighbors": 15}, + "data": np.arange(18, dtype=np.float32).reshape(6, 3) / 4.0, + }, +] + + +def source_tables() -> list[pa.Table]: + """The three v2-shaped tables the prepare pipeline would hand ``write_bundle``.""" + from protlabel import Prediction + + processor = BaseProcessor({}, {}) + frame = pd.DataFrame({"identifier": PROTEIN_IDS, **ANNOTATION_CELLS}) + annotations = processor._create_protein_annotations_table(frame) + + annotations = add_overlay_columns( + annotations, + "kingdom", + [ + Prediction( + query_id=q, + label=lab, + source_id=s, + distance=d, + reliability=r, + k=1, + metric="euclidean", + ) + for q, lab, r, d, s in PREDICTIONS + ], + identifiers=PROTEIN_IDS, + ) + # append_column/drop_columns are not guaranteed to carry schema metadata. + annotations = stamp_format_version(annotations) + + return [ + annotations, + processor._create_projections_metadata_table(PROJECTIONS), + processor._create_projections_data_table(PROJECTIONS, PROTEIN_IDS), + ] + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(message)s") + write_bundle(source_tables(), FIXTURE_PATH) + print(f"{FIXTURE_PATH} ({FIXTURE_PATH.stat().st_size} bytes)") + + +if __name__ == "__main__": + main() diff --git a/apps/protspace/tests/test_bundle_v3_fixture.py b/apps/protspace/tests/test_bundle_v3_fixture.py new file mode 100644 index 00000000..14fc5cf1 --- /dev/null +++ b/apps/protspace/tests/test_bundle_v3_fixture.py @@ -0,0 +1,198 @@ +"""The golden v3 bundle both languages read. + +``packages/core/src/components/data-loader/utils/__fixtures__/v3-sample.parquetbundle`` +is committed binary: Python writes it (``scripts/generate_v3_fixture.py``) and +vitest reads it, so it is the only place the two v3 implementations meet. These +tests keep the committed bytes and the generator from drifting apart, and pin the +exact cells the browser side asserts on -- a fixture nobody reads back in Python +is a fixture that silently rots. + +Regenerate with ``uv run python scripts/generate_v3_fixture.py``. +""" + +import importlib.util +import io +import json +import sys +from pathlib import Path + +import pyarrow.parquet as pq +import pytest + +from protspace.data.io.bundle import ( + PARQUET_BUNDLE_DELIMITER, + read_settings_from_bundle, + read_tables, +) +from protspace.data.io.bundle_v3 import MANIFEST_KEY + +FIXTURE = ( + Path(__file__).resolve().parents[3] + / "packages" + / "core" + / "src" + / "components" + / "data-loader" + / "utils" + / "__fixtures__" + / "v3-sample.parquetbundle" +) + +_SPEC = importlib.util.spec_from_file_location( + "generate_v3_fixture", + Path(__file__).parent.parent / "scripts" / "generate_v3_fixture.py", +) +generator = importlib.util.module_from_spec(_SPEC) +sys.modules["generate_v3_fixture"] = generator +_SPEC.loader.exec_module(generator) + +pytestmark = pytest.mark.skipif( + not FIXTURE.exists(), reason="browser fixtures not checked out" +) + + +@pytest.fixture(scope="module") +def parts() -> list[bytes]: + return FIXTURE.read_bytes().split(PARQUET_BUNDLE_DELIMITER) + + +@pytest.fixture(scope="module") +def tables(): + return read_tables(FIXTURE) + + +def test_fixture_is_a_six_part_v3_container(parts): + assert len(parts) == 6 + # No settings, no statistics -- the two zero-byte slots are what keeps the + # payloads part at index five, where the browser reads it positionally. + assert parts[3] == b"" and parts[4] == b"" + assert read_settings_from_bundle(FIXTURE) is None + + footer = pq.read_metadata(io.BytesIO(parts[0])).metadata + assert footer[b"protspace_format_version"] == b"3" + assert pq.read_table(io.BytesIO(parts[5])).column_names == ["name", "data"] + + +def test_manifest_declares_every_kind_the_reader_dispatches_on(parts): + manifest = json.loads(pq.read_metadata(io.BytesIO(parts[0])).metadata[MANIFEST_KEY]) + + assert manifest["idColumn"] == "protein_id" + assert manifest["projections"] == [ + {"name": "pca2", "dimension": 2}, + {"name": "umap3", "dimension": 3}, + ] + assert manifest["columns"] == { + "cath": {"kind": "multi", "sourceType": "string", "scores": True}, + "go_bp": {"kind": "multi", "sourceType": "string", "evidence": True}, + "pfam": {"kind": "multi", "sourceType": "string", "scores": True}, + "kingdom": {"kind": "categorical", "sourceType": "string"}, + "predicted_tm": {"kind": "categorical", "sourceType": "string"}, + "length": {"kind": "numeric", "numericType": "int", "sourceType": "string"}, + "hydrophobicity": { + "kind": "numeric", + "numericType": "float", + "sourceType": "string", + }, + "kingdom__pred_value": {"kind": "categorical", "sourceType": "string"}, + "kingdom__pred_confidence": { + "kind": "numeric", + "numericType": "float", + "sourceType": "float", + }, + "kingdom__pred_source": {"kind": "categorical", "sourceType": "string"}, + } + + +def test_part_one_is_the_physical_v3_schema(parts): + schema = pq.read_schema(io.BytesIO(parts[0])) + + # Multi columns become a hit count; categoricals a code; numerics a double. + assert schema.names == [ + "protein_id", + "cath__count", + "go_bp__count", + "pfam__count", + "kingdom", + "predicted_tm", + "length", + "hydrophobicity", + "kingdom__pred_value", + "kingdom__pred_confidence", + "kingdom__pred_source", + ] + # hyparquet only hands back typed arrays for REQUIRED flat columns. + assert all(not field.nullable for field in schema) + + +def test_read_tables_gives_back_the_v2_cells_the_generator_started_from(tables): + annotations, _metadata, _data = tables + cells = generator.ANNOTATION_CELLS + + assert annotations.column("protein_id").to_pylist() == generator.PROTEIN_IDS + # P1/P2's cath and go_bp are the v2 golden fixture's own cells: the + # percent-encoded ';' inside a CATH name survives the hit split, an unnamed + # bare code stays unnamed, and the '|IDA' evidence suffix round-trips. + for column in ("cath", "go_bp", "pfam", "kingdom", "predicted_tm", "length"): + assert annotations.column(column).to_pylist() == cells[column], column + + # A score only float64 can carry: 1e-200 flushes to zero in float32 and + # 123456789 re-spells as 1.2345679e+08. + assert "|1e-200" in annotations.column("cath")[4].as_py() + assert annotations.column("cath")[2].as_py().endswith("|123456789") + + # Documented non-identities: a float numeric column comes back in its + # shortest round-trip spelling, and a null string cell as "". + assert annotations.column("hydrophobicity").to_pylist() == [ + "0.5", + "-1.25", + "", + "3.0", + "0.001", + "42.0", + ] + assert annotations.column("kingdom__pred_source").to_pylist() == [ + "", + "Q9XYZ1", + "", + "P0A7B8", + "A0A123", + "", + ] + # sourceType restores the EAT confidence's float32, missing as NaN. + confidence = annotations.column("kingdom__pred_confidence") + assert confidence.type == "float" + assert [confidence[i].as_py() for i in (1, 3, 4)] == [0.875, 0.5, 0.25] + + +def test_read_tables_rebuilds_both_projections_in_protein_order(tables): + _annotations, metadata, data = tables + + assert metadata.column("projection_name").to_pylist() == ["pca2", "umap3"] + assert metadata.column("dimensions").to_pylist() == [2, 3] + + frame = data.to_pandas() + pca2 = frame[frame.projection_name == "pca2"] + assert pca2.identifier.tolist() == generator.PROTEIN_IDS + # P1 and P2 sit where the v2 golden fixture puts them. + assert pca2[["x", "y"]].values.tolist()[:2] == [[0.0, 0.0], [1.0, 1.0]] + assert pca2.z.isna().all() # 2D: no z + + umap3 = frame[frame.projection_name == "umap3"] + assert umap3[["x", "y", "z"]].values.tolist()[0] == [0.0, 0.25, 0.5] + + +def test_committed_fixture_still_matches_its_generator(tmp_path): + """The committed bytes are what ``scripts/generate_v3_fixture.py`` writes today. + + Compared decoded rather than byte-for-byte: parquet embeds the writer + version, so bytes would fail on every pyarrow bump for no reason. + """ + from protspace.data.io.bundle import write_bundle + + regenerated = tmp_path / "v3-sample.parquetbundle" + write_bundle(generator.source_tables(), regenerated) + + for fresh, committed in zip( + read_tables(regenerated), read_tables(FIXTURE), strict=True + ): + assert fresh.equals(committed) diff --git a/packages/core/src/components/data-loader/utils/__fixtures__/v3-sample.parquetbundle b/packages/core/src/components/data-loader/utils/__fixtures__/v3-sample.parquetbundle new file mode 100644 index 0000000000000000000000000000000000000000..3bf7cbc274209a48ecea20b5ea27fca9f4becb3f GIT binary patch literal 12809 zcmdT~TWlLwdY-WHq{xW0M9aL@S!QN<8)`V@ z@*;`KdSmoqQ3QS{g1)q9i@dfdiegdpWz(!5i=sdmivqo&t$`Nk0`0@@wm<_kK>Ph? zE*f6s*of9Z2b|%|IsfgP^Pm6x|9{Sokn*`W*OqI2(KU&eao6}P$8oP;Ifv(tZ)}|N z;|`3CdskS=DtlfBq}R3BhvK(D(N!C0MQRz8+}B*95=>0zO_Df!O4B` z(U{Y*-krKU&A-0%h+%TSvCfcoMj${JT{xzW`+ddXy^QIMK_ zMK3hVV!f~A!oH$aijBH~(rlHAirgs`3mp_)1l@79q*Xdq{o}=Iuh?ods|`t!l}=Ab zdf6aRC=hnxKvcUW9o*AjBR90NQY>jQs!8N4{nhP8w}mSAg zA{hNPYA7w?e?B&v%c-Mm@~bZV#rOcK!EdAnSPZ{^!G`7jd#8JH!Rh`rCpWY08E=YO zQEYpDHu2%Sq7_R=P<*vp*V^;9=6?0Vc_>y9Z}TlttCZ#!=jJuBUP6J^RaJboJ7TMY zPcM6KH1T?!y_V?JcmI8S==2JVc*!8!Po=Ds)MBB(Oc$%r={1Snd{<~8L0bOQc<7X zUxaW*OXtbZ40AL~&z2W6VZ+R<;jZuZ1h zt8qw@9h6#arJ+%&e|5>b5KEQcg zC$$HGLS599a;e?%d@>G8Sqw8PIe)9qKp3_T=H)|ua^qjOp!4;v+EL~g=k>LDkC`=? z&lcGzl*)}3wY1afmhLak>$CD{kxK`QlU4b?rt(z^8$_<@`8W^|XL!O)NFK}ay5h+j>utTx0a&@(#1?u*f|8pHWiozhR@ z#A&n9IjyN9h0_bs)_%geO4*1$g>Y=h8nkp`!a(t`)a$F#d{Y+v=;fkPFKPO0^50*y zYP)r@8R)ML+Z?Db!cD9Gz1^7pjVD#bDjR|jiqSjp@O^x9dHH6=`VPgY{zj~>?(h18?PA8)jnx{8 zycItvM$=764L0-PSUXlv9!Oj1j-(}&+h}CYAMuLe*!5B`-V!sb`-CH9)9qq(U)c`F zkEKAmm-DCZma-w0(F>${fYV~ip(6W}2XZ}q9IG|GvX)lMcMg@jKUmIb0;4VZ)h>8m zkJWuXX)D##ae_X%o>bFY>S6v4>a?U+t^=+|vZeiECZ$TKx6_oP`Kkoivyzd?!Vel zlu2U@lGd1DJiJft+jtOT*(haFJt+`sNPf`rtKQOJpMC|OeEsMkm)-s9eyMkwXve2A(H zNZ;(M1j*dRoKR_OE#0f<+Y-inBww?R#(6}0WWTP*&{NDgYVWRcTZz|lS?Gd4NaItL zXufI5gJL#b%coPTauo6@@#ErF9PqMI&gwcO`I{Z5NJ{K_0&>~m(6*|)!P>bm0$0+Kptqvjt|f!ntQV%?E{sM9eHQe}flw85 zx?HR!%TgxN6*GBN4v?wZNRa7*5yHhx7c(^pvOh3hgZTM+um^l^HUNB`A0Y~1K}Zlj z{>g^nBHoPt5fOykXn#sAMS?ZSe*}G|b~h@SKwGhFMmXNf(OBO>+c9=Ve^ewb-Il_K z6|x(-Y}`wD3FBJ|!0sr5e>*HFGUlXAwr3*}%-#XKY4>0U8Sg0BEuE)Sk93$Vib7wa zhs3WgX3}lU=@EQGRp3`;U(jpVM97jxA0`j-7+>NI89VZpS}vM4E*w9|)oF}r?w~%C z_U1Av=>8_;C7VsU8^pN5u0p1mn!cM-yel!y*DpKSY$asZ5AY6QPN!7hmc@)tKVMf> zDVi!fJqzKmf#dIhd>E@%6<7~j3yMj{8tio0W_f5UTygMwVZMskoDq=ld zE7IH#$6>cqRnYL}Xsu)Vbp%^XV}8QZSl>qK9mc7i%XnG;$lC*ANe)-g5Gy)o_D%lQZAGffsY_%2J{2#UUpCVs%H4mRrj#go_)S8%0tuj8p zsQybs)^NO9+#7P;wL6(>QDvWOVHB+S zmJXJg?nZ)!-%abLAVhX(z%w|Coc?u)g*2WK3-GY!NAIJa-!1qiR$Yd{Nx&9xNeac8U7u6W%6un4#=}lb%XM- z7%_4p5Be9%P=PmT`gcGwl|ID$J`~6T>sJsv8T7j&`=7c*;{;WB`)$rM z*A!dd?v^^&oxNk_OhP*<6in@?0#Zv^xCy+_}jg^xmVcO+>Rw; zdz&e}jCM_zXamoI7h~7-m$I~@=HMN8Ql4N8|12W_kVO}EL2a;Zw!`Y+`Ej;idLE#3 z0kY^S_Cfkpo$55IkvD79n+>g`btp;izyI=>Yh5d7C8bi88Z9dGt=?AkM&4IfZ10{HsraQl|giH;64@2?x+i`?n;6~aNr(YAeN z729^thW4qfc`(eguWEqhv+dapaW&GK5V?Ox?xg(d^LuVP4gKlShm*AAEUz%rZmz-4 zo*s=-I@DMm7!|fYVr4(1Wdr}@Dc}vocw*(V(xV9J(km=&WTrv&klIXt!c~kTG%lL% zw-FkKSa@rs7OJvF=^o+G_&(564J_n4h6{U0u4N+YD^fj{ah8rn=*@WNv%_J$mDU&v zNo{=3I{ZV3K&~J{4VI3g@Hd@KMYs1VNKN(2F82_=MmqgrE5Za8>PC^0QjJsz3rU8< zYDzwVTlx%X6dIQi=?Giu3@NG8(j-V%7dLW7_%$4!Q93qR8=2}kB}D?g zs0Yg;QiVwEof4J}!91_f(F{dS>$ARkGjy_tt7nli(NlO4FQra%TH0&Op&Kl;t0$3$ zk9eg(%t)sl0Va+)y3bNNNX6kD+aQ!a561_?KOaKGZX!%6a-W5(r2c9VU9biA>SAm8KsaJK7>LkRyY_gu<*r9B?F#?2X)iE-__k^7E{>5 zI~JKmZTdi_F1ANj)pX}Wf*xKBy7REJf`AlC2lJkVz%t|?>CWqsCwy?`&O^=-?z|yi zoHflfLly1acy>O6GZZ{q<@ z(9NxL+=E9P_bL7#{Se;t>7X9~?7=%aD5oFl8%1>;R3G8K42tWxn>y%9zX11iS|q|_ zLyR2&KR6AT1JI|qIRJTae}Mbf*Woa$HrM_l=W^l<+yx3q3WcM-FC46+IwuQKtV;89 zp!QflF?WxAa;NT-n?&I%PrB3Yw9SjxefNI2o(1TWC#AhuD)h+z@o)Mp<6T%q_KRJlc1hmOd8H~_j&cE41s$o?~Yy*|ryGf0N6=Co& zVXWQu$$v^(sE?lBj}_|ooNjkXJY%$;CLM6kG93B~CmF*lnA;$t^?@wLk?{L%8uRC!@g`d5`PUO0hE{yz(R-TdntVzbfgG&+hrhgelEwcE|9A8eN!Djk-jfuSVKCEQoW zViP=HCBpA`IR1CCl)!ln_fEu%%QtuRMyLLk@6CxzoPXkiCop$a{EqY{|DSx}$^}48 zaTEW;^V0?X_N$j5-j!*`;#Cy?;LrZ%@BaBa|L`{N1uqBxFLeGLE5R7g{b0wqD756NSQ0n8`#9*;=!pCS{@?fqd!$k?;?|HO+Dpt;OR92|6!>{3vMFN= z+_`I0w~`z$K<>@GmjJgL%;xgG9j|aBbW`vOK7WAUbmEY+%W;*MBO=-i-@f(gYjghV z!5HRic8CA>SKnED0d>@3y0PA@+149}etO(N=VeY3L!8G!hfg`T7YU(7TSADz!8Ci_ ze}2LJlk>I}WX=6gW4065MskLcVelVcbpQ2*g`b|BU0!p4>NF7CKc04bk(>GKg<-4y zm?z7%e4WgMX(#CL&9ho{GW5)T>(6A(uxi34GkvY^A j38TF`#x=?|2N@}RgO9mD3BuD@y+y;=eU2$ZHtPQXJvl5Z literal 0 HcmV?d00001 From 0546d4936fab1e683c1fba773a8d80feeebe4c4f Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 03:01:24 +0200 Subject: [PATCH 27/31] perf(bundle): stop paying a v3 core decode for a settings read transfer, inspect and the toxprot script take read_tables directly. --- .../scripts/generate_toxprot_demo.py | 14 ++--- apps/protspace/scripts/inspect_bundle.py | 13 ++-- apps/protspace/src/protspace/cli/bundle.py | 4 ++ apps/protspace/src/protspace/cli/transfer.py | 13 ++-- .../src/protspace/data/io/__init__.py | 2 + .../protspace/src/protspace/data/io/bundle.py | 11 ++++ .../protspace/utils/add_annotation_style.py | 14 +++-- apps/protspace/tests/test_bundle_settings.py | 59 +++++++++++++++++++ 8 files changed, 102 insertions(+), 28 deletions(-) diff --git a/apps/protspace/scripts/generate_toxprot_demo.py b/apps/protspace/scripts/generate_toxprot_demo.py index 0a6cdfbc..50dcaeee 100644 --- a/apps/protspace/scripts/generate_toxprot_demo.py +++ b/apps/protspace/scripts/generate_toxprot_demo.py @@ -24,7 +24,6 @@ from pathlib import Path import pyarrow as pa -import pyarrow.parquet as pq import requests logger = logging.getLogger(__name__) @@ -301,15 +300,16 @@ def postprocess_bundle( """Patch the bundle: mature lengths, column drop+reorder, restyled top-9 categories, and the original ``protein_families`` settings. """ - from protspace.data.io.bundle import read_bundle, write_bundle + from protspace.data.io.bundle import ( + read_settings_from_bundle, + read_tables, + write_bundle, + ) if not source_settings_bundle.exists(): raise SystemExit(f"Source settings bundle not found: {source_settings_bundle}") - parts, _ = read_bundle(bundle_path) - annotations = pq.read_table(io.BytesIO(parts[0])) - metadata = pq.read_table(io.BytesIO(parts[1])) - data = pq.read_table(io.BytesIO(parts[2])) + annotations, metadata, data = read_tables(bundle_path) # Map by protein_id (not positional) — bundle row order is not guaranteed # to match FASTA order after EmbeddingSet merging and dedup in the prepare @@ -331,7 +331,7 @@ def postprocess_bundle( annotations = _drop_and_reorder_columns(annotations) - _, source_settings = read_bundle(source_settings_bundle) + source_settings = read_settings_from_bundle(source_settings_bundle) if source_settings is None: raise SystemExit( f"Source settings bundle has no settings part: {source_settings_bundle}" diff --git a/apps/protspace/scripts/inspect_bundle.py b/apps/protspace/scripts/inspect_bundle.py index c601272d..2791b844 100644 --- a/apps/protspace/scripts/inspect_bundle.py +++ b/apps/protspace/scripts/inspect_bundle.py @@ -2,12 +2,9 @@ """Inspect a .parquetbundle: per-table rows/cols/schema, sample, and settings.""" import argparse -import io from pathlib import Path -import pyarrow.parquet as pq - -from protspace.data.io.bundle import read_bundle +from protspace.data.io.bundle import read_settings_from_bundle, read_tables TABLE_NAMES = ["selected_annotations", "projections_metadata", "projections_data"] @@ -20,10 +17,12 @@ def main(): ) args = parser.parse_args() - parts, settings = read_bundle(Path(args.bundle)) + # read_tables decodes a v3 core back to the v2 shape, so a v3 bundle + # inspects as the three tables it logically holds, not as its parts. + tables = read_tables(Path(args.bundle)) + settings = read_settings_from_bundle(Path(args.bundle)) - for name, blob in zip(TABLE_NAMES, parts, strict=True): - table = pq.read_table(io.BytesIO(blob)) + for name, table in zip(TABLE_NAMES, tables, strict=True): print(f"== {name}: {table.num_rows} rows × {table.num_columns} cols ==") print(table.schema) if args.sample_rows > 0 and table.num_rows > 0: diff --git a/apps/protspace/src/protspace/cli/bundle.py b/apps/protspace/src/protspace/cli/bundle.py index 2a1b6bd0..81e3ba4e 100644 --- a/apps/protspace/src/protspace/cli/bundle.py +++ b/apps/protspace/src/protspace/cli/bundle.py @@ -98,6 +98,10 @@ def bundle( # same-version annotate/prepare pipeline (i.e. already percent-encoded). # We don't inspect its contents, so it's unconditionally stamped as v2 -- # there is currently no other producer of this parquet to distrust. + # The stamp declares the *cell grammar*, not the container: write_bundle + # emits a v3 container either way, and without the v2 stamp its encoder + # would take the table for v1 and migrate it, double-escaping every + # reserved character. annotations_table = stamp_format_version(annotations_table) statistics_table = ( diff --git a/apps/protspace/src/protspace/cli/transfer.py b/apps/protspace/src/protspace/cli/transfer.py index 1ebd7eb6..74e67ee6 100644 --- a/apps/protspace/src/protspace/cli/transfer.py +++ b/apps/protspace/src/protspace/cli/transfer.py @@ -234,16 +234,12 @@ def transfer( """ setup_logging(verbose) - import io - - import pyarrow.parquet as pq - from protspace.analysis.classification import Rule from protspace.data.annotations.encoding import ( migrate_legacy_annotation_table, read_format_version, ) - from protspace.data.io.bundle import read_bundle, replace_annotations_in_bundle + from protspace.data.io.bundle import read_tables, replace_annotations_in_bundle from protspace.data.loaders import load_h5, split_h5_spec def _parse_where(items: list[str] | None) -> list[tuple[str, str]]: @@ -275,9 +271,10 @@ def _parse_where(items: list[str] | None) -> list[tuple[str, str]]: emb_set = load_h5([h5_path], name_override=name_override) emb_map = {header: emb_set.data[i] for i, header in enumerate(emb_set.headers)} - # Read the annotations part of the bundle. - parts, _settings = read_bundle(bundle) - annotations = pq.read_table(io.BytesIO(parts[0])) + # Read the annotations part of the bundle. read_tables hands back the v2 + # shape for a v3 container too, without the parquet round trip read_bundle + # would do on the way out of the decoder. + annotations, _metadata, _projections = read_tables(bundle) input_format_version = read_format_version(annotations) # Real bundles name the id column "protein_id"; run_transfer works on "identifier". diff --git a/apps/protspace/src/protspace/data/io/__init__.py b/apps/protspace/src/protspace/data/io/__init__.py index 1cf6124c..9b30ac18 100644 --- a/apps/protspace/src/protspace/data/io/__init__.py +++ b/apps/protspace/src/protspace/data/io/__init__.py @@ -12,6 +12,7 @@ create_settings_parquet, extract_bundle_to_dir, read_bundle, + read_settings_from_bundle, read_settings_from_bytes, read_settings_from_file, read_statistics_from_bundle, @@ -32,6 +33,7 @@ "PARQUET_BUNDLE_DELIMITER", "extract_bundle_to_dir", "read_bundle", + "read_settings_from_bundle", "read_statistics_from_bundle", "read_tables", "write_bundle", diff --git a/apps/protspace/src/protspace/data/io/bundle.py b/apps/protspace/src/protspace/data/io/bundle.py index d86ee5a7..64ef8ab0 100644 --- a/apps/protspace/src/protspace/data/io/bundle.py +++ b/apps/protspace/src/protspace/data/io/bundle.py @@ -253,6 +253,17 @@ def read_bundle(bundle_path: Path) -> tuple[list[bytes], dict | None]: return core, settings +def read_settings_from_bundle(bundle_path: Path) -> dict | None: + """Return the parsed settings (fourth part), or None if absent. + + The settings-only read: unlike :func:`read_bundle` it never touches the + core, so ``protspace style`` no longer pays a full v3 decode plus + re-serialization of every annotation column just to look at a JSON blob. + """ + settings = _parse_bundle(bundle_path)[1] + return read_settings_from_bytes(settings) if settings else None + + def read_statistics_from_bundle(bundle_path: Path) -> bytes | None: """Return the raw statistics parquet bytes (fifth part), or None if absent.""" return _parse_bundle(bundle_path)[2] diff --git a/apps/protspace/src/protspace/utils/add_annotation_style.py b/apps/protspace/src/protspace/utils/add_annotation_style.py index 0bbc4df3..9cbc98ca 100644 --- a/apps/protspace/src/protspace/utils/add_annotation_style.py +++ b/apps/protspace/src/protspace/utils/add_annotation_style.py @@ -350,7 +350,7 @@ def add_annotation_styles_bundle( """ from protspace.data.io.bundle import ( extract_bundle_to_dir, - read_bundle, + read_settings_from_bundle, replace_settings_in_bundle, ) from protspace.data.io.settings_converter import visualization_state_to_settings @@ -359,8 +359,10 @@ def add_annotation_styles_bundle( temp_dir = extract_bundle_to_dir(Path(bundle_file)) reader = ArrowReader(Path(temp_dir)) - # Read existing settings from the bundle (if any) to preserve extra fields - _, existing_settings = read_bundle(Path(bundle_file)) + # Read existing settings from the bundle (if any) to preserve extra fields. + # Settings-only: extract_bundle_to_dir above already paid the one v3 decode + # this command needs. + existing_settings = read_settings_from_bundle(Path(bundle_file)) # Collect settings-level overrides from the styles input style_overrides: dict[str, dict] = {} @@ -435,11 +437,11 @@ def dump_settings(input_file: str) -> None: data_format = detect_data_format(input_file) if data_format == "parquetbundle": - from protspace.data.io.bundle import read_bundle + from protspace.data.io.bundle import read_settings_from_bundle - _, settings = read_bundle(Path(input_file)) + settings = read_settings_from_bundle(Path(input_file)) if settings is None: - print("No settings found in bundle (3-part bundle).") + print("No settings found in bundle.") else: print(json.dumps(settings, indent=2)) elif data_format == "parquet": diff --git a/apps/protspace/tests/test_bundle_settings.py b/apps/protspace/tests/test_bundle_settings.py index dc305a50..66ea4072 100644 --- a/apps/protspace/tests/test_bundle_settings.py +++ b/apps/protspace/tests/test_bundle_settings.py @@ -34,3 +34,62 @@ def test_list_values(self): data = create_settings_parquet(original) result = read_settings_from_bytes(data) assert result == original + + +class TestSettingsOnlyRead: + """``read_settings_from_bundle`` must agree with ``read_bundle``'s second value. + + ``protspace style`` reads settings and nothing else; on a v3 container + ``read_bundle`` would decode and re-serialize every annotation column first, + so the settings-only reader exists purely to skip that. It earns its place + only while it stays byte-for-byte equivalent to the path it replaced. + """ + + def test_matches_read_bundle_on_v3_legacy_and_absent(self, tmp_path): + import io + + import pyarrow as pa + import pyarrow.parquet as pq + + from protspace.data.io.bundle import ( + PARQUET_BUNDLE_DELIMITER, + read_bundle, + read_settings_from_bundle, + write_bundle, + ) + + settings = {"family": {"categories": {"kinase": {"color": "#FF0000"}}}} + tables = [ + pa.table({"protein_id": ["p0", "p1"], "family": ["kinase", ""]}), + pa.table({"projection_name": ["pca2"], "dimensions": [2]}), + pa.table( + { + "projection_name": ["pca2"] * 2, + "identifier": ["p0", "p1"], + "x": [0.0, 1.0], + "y": [2.0, 3.0], + } + ), + ] + + v3 = tmp_path / "v3.parquetbundle" + write_bundle(tables, v3, settings=settings) + + bare = tmp_path / "bare.parquetbundle" + write_bundle(tables, bare) + + def serialized(table): + buf = io.BytesIO() + pq.write_table(table, buf) + return buf.getvalue() + + legacy = tmp_path / "legacy.parquetbundle" + legacy.write_bytes( + PARQUET_BUNDLE_DELIMITER.join( + [*(serialized(t) for t in tables), create_settings_parquet(settings)] + ) + ) + + for path, expected in ((v3, settings), (legacy, settings), (bare, None)): + assert read_settings_from_bundle(path) == expected, path.name + assert read_settings_from_bundle(path) == read_bundle(path)[1], path.name From c4932f06d5f0669f9190dd9d517c2cfdee07df06 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 03:12:40 +0200 Subject: [PATCH 28/31] test(core): prove the v3 encoder and reader agree on the golden bundle Also pins the v3-to-v2 export round trip and the shared v2 anchors. --- .../data-loader/utils/v3-equivalence.test.ts | 459 ++++++++++++++++++ 1 file changed, 459 insertions(+) create mode 100644 packages/core/src/components/data-loader/utils/v3-equivalence.test.ts diff --git a/packages/core/src/components/data-loader/utils/v3-equivalence.test.ts b/packages/core/src/components/data-loader/utils/v3-equivalence.test.ts new file mode 100644 index 00000000..f1cc660e --- /dev/null +++ b/packages/core/src/components/data-loader/utils/v3-equivalence.test.ts @@ -0,0 +1,459 @@ +import { describe, it, expect } from 'vitest'; +import { readFileSync } from 'node:fs'; +import { + createParquetBundle, + getProteinAnnotationIndices, + getProteinEvidence, + getProteinScores, + isCsrAnnotationData, + isMultilabelAnnotationData, + isNAValue, + NA_DEFAULT_COLOR, + NA_VALUE, + type VisualizationData, +} from '@protspace/utils'; +import { decodeParquetBundle, extractRowsFromParquetBundle } from './bundle'; +import { convertParquetToVisualizationDataOptimized } from './conversion'; + +/** + * The cross-language contract for parquetbundle v3. + * + * `__fixtures__/v3-sample.parquetbundle` is written by the real Python encoder + * (`apps/protspace/scripts/generate_v3_fixture.py` -> `bundle_v3.encode_v3`) and read + * here by the real browser reader. `bundle-v3.test.ts` covers the reader against + * hand-synthesised parts, where a corrupt manifest or a lying hit count is reachable; + * this suite covers the half that no synthesised part can: that the two independently + * written implementations agree on the same bytes, and that a v3-loaded dataset behaves + * like a v2-loaded one everywhere downstream. + * + * The fixture is a deliberate SUPERSET of `v2-sample.parquetbundle` (2 rows, 2 columns, + * which cannot carry an interior zero-hit row, a numeric column or a 3D projection), so + * the two are compared only on the values they genuinely share: `protein_ids[0..1]`, + * P1/P2's `cath` and `go_bp`, and the `pca2` coordinates. + */ + +const PROTEIN_IDS = ['P1', 'P2', 'P3', 'P4', 'P5', 'P6'] as const; + +/** Every categorical column of the fixture, in part-1 order. */ +const CATEGORICAL = ['cath', 'go_bp', 'pfam', 'kingdom', 'predicted_tm'] as const; + +/** + * The label whose own text contains the `;` the v2 cell grammar reserves as the hit + * separator, so it travels percent-encoded and must come back decoded. Byte-identical + * to the one in `v2-sample.parquetbundle`. + */ +const CATH_ENCODED_SEMICOLON = 'G3DSA:1.10.10.10 (Ribosomal Protein L15; Chain: K; domain 2)'; + +/** + * Rows whose only "hit" is the synthetic `__NA__` the reader inserts for an empty CSR + * row, per column that carries a score or evidence payload. These are exactly the rows + * where CSR and nested storage legitimately disagree (see the round-trip suite). + */ +const NA_ONLY_ROWS: Readonly> = { + cath: ['P4'], + go_bp: ['P2', 'P5'], + pfam: ['P1', 'P3', 'P6'], +}; + +function fixture(name: string): ArrayBuffer { + const file = readFileSync(new URL(`./__fixtures__/${name}`, import.meta.url)); + return file.buffer.slice(file.byteOffset, file.byteOffset + file.byteLength) as ArrayBuffer; +} + +/** The v3 fixture through the real version-sniffing entry point. */ +const loadV3 = () => decodeParquetBundle(fixture('v3-sample.parquetbundle')); + +/** Any v2 container through the legacy row-object path, which is what v3 must match. */ +const loadLegacy = async (buffer: ArrayBuffer): Promise => + convertParquetToVisualizationDataOptimized(await extractRowsFromParquetBundle(buffer)); + +/** One protein's view of one column, read only through the public accessors. */ +interface Hits { + labels: (string | null)[]; + scores: (number[] | null)[]; + evidence: (string | null)[]; +} + +function hitsOf(data: VisualizationData, key: string, proteinIndex: number): Hits { + const annotation = data.annotations[key]; + return { + labels: getProteinAnnotationIndices(data.annotation_data[key], proteinIndex).map( + (valueIndex) => annotation.values[valueIndex], + ), + scores: getProteinScores(data, proteinIndex, key), + evidence: getProteinEvidence(data, proteinIndex, key), + }; +} + +const hitsByProtein = (data: VisualizationData, key: string): Record => + Object.fromEntries(data.protein_ids.map((id, index) => [id, hitsOf(data, key, index)])); + +describe('v3 golden fixture: the Python encoder and the browser reader agree', () => { + it('reads the six-part container with both empty slots and no settings or statistics', async () => { + const { data, settings } = await loadV3(); + + expect(data.protein_ids).toEqual([...PROTEIN_IDS]); + // Parts 4 and 5 are the zero-byte slots that keep the payloads part at position six. + expect(settings).toBeNull(); + expect(data.statistics).toBeUndefined(); + }); + + it('exposes exactly the declared annotations, with the EAT companion trio consumed', async () => { + const { data } = await loadV3(); + + // `kingdom__pred_value/__pred_confidence/__pred_source` are declared in the manifest + // and physically present in part 1; `normalizeEatCompanionColumns` must consume them + // so they never become three junk legend columns. + expect(Object.keys(data.annotations)).toEqual([ + 'cath', + 'go_bp', + 'pfam', + 'kingdom', + 'predicted_tm', + 'length', + 'hydrophobicity', + ]); + // Every prediction targets a protein whose curated `kingdom` is present, so the + // overlay yields no cells at all - and must not invent an empty record either. + expect(data.annotation_predicted).toBeUndefined(); + }); + + it('stores multi-valued columns as CSR and single-valued ones as flat codes', async () => { + const { data } = await loadV3(); + + for (const key of ['cath', 'go_bp', 'pfam'] as const) { + const storage = data.annotation_data[key]; + expect(isCsrAnnotationData(storage), key).toBe(true); + expect(isMultilabelAnnotationData(storage), key).toBe(true); + expect((storage as { length: number }).length).toBe(PROTEIN_IDS.length); + } + for (const key of ['kingdom', 'predicted_tm'] as const) { + expect(data.annotation_data[key], key).toBeInstanceOf(Int32Array); + } + }); + + it('decodes the scored multi column, keeping 1e-200 and 123456789 exact', async () => { + const { data } = await loadV3(); + + // Frequency-sorted, ties by first occurrence; every label appears twice here, so + // this pins first-occurrence order. `__NA__` is appended last, never sorted in. + expect(data.annotations.cath).toEqual({ + kind: 'categorical', + values: [CATH_ENCODED_SEMICOLON, 'G3DSA:6.20.10.10', '6.20.10.10', NA_VALUE], + colors: ['#F3C300', '#875692', '#F38400', NA_DEFAULT_COLOR], + shapes: ['circle', 'circle', 'circle', 'circle'], + }); + // The percent-encoded ';' came back as one label, not three fragments. + expect(data.annotations.cath.values.some((value) => value?.includes('%3B'))).toBe(false); + + expect(hitsByProtein(data, 'cath')).toEqual({ + P1: { + labels: [CATH_ENCODED_SEMICOLON, 'G3DSA:6.20.10.10'], + scores: [[50.2], [60.5]], + evidence: [], + }, + P2: { labels: ['6.20.10.10'], scores: [null], evidence: [] }, + P3: { labels: ['G3DSA:6.20.10.10'], scores: [[123456789]], evidence: [] }, + P4: { labels: [NA_VALUE], scores: [null], evidence: [] }, + P5: { labels: [CATH_ENCODED_SEMICOLON], scores: [[1e-200]], evidence: [] }, + P6: { labels: ['6.20.10.10'], scores: [null], evidence: [] }, + }); + + // The two values the wire format is float64 for. float32 flushes 1e-200 to zero and + // re-spells 123456789 as 123456792, so `toEqual` above would already fail - these + // two say why, and fail loudly if someone widens the comparison instead. + const [eValue] = getProteinScores(data, PROTEIN_IDS.indexOf('P5'), 'cath')[0]!; + expect(eValue).toBe(1e-200); + expect(Math.fround(eValue)).toBe(0); + const [large] = getProteinScores(data, PROTEIN_IDS.indexOf('P3'), 'cath')[0]!; + expect(large).toBe(123456789); + expect(Math.fround(large)).not.toBe(large); + }); + + it('decodes the evidenced multi column against the global evidence dictionary', async () => { + const { data } = await loadV3(); + + expect(data.annotations.go_bp).toEqual({ + kind: 'categorical', + // 'apoptotic process' has 3 hits, 'protein folding' 2. + values: ['apoptotic process', 'protein folding', NA_VALUE], + colors: ['#F3C300', '#875692', NA_DEFAULT_COLOR], + shapes: ['circle', 'circle', 'circle'], + }); + + expect(hitsByProtein(data, 'go_bp')).toEqual({ + P1: { labels: ['apoptotic process'], scores: [], evidence: ['IDA'] }, + P2: { labels: [NA_VALUE], scores: [], evidence: [null] }, + P3: { + labels: ['apoptotic process', 'protein folding'], + scores: [], + // Both evidence spellings the grammar accepts, on one protein. + evidence: ['IDA', 'ECO:0000269'], + }, + P4: { labels: ['protein folding'], scores: [], evidence: ['IEA'] }, + P5: { labels: [NA_VALUE], scores: [], evidence: [null] }, + P6: { labels: ['apoptotic process'], scores: [], evidence: ['EXP'] }, + }); + }); + + it('prefix-sums CSR rows correctly with zero hits first, interior and last', async () => { + const { data } = await loadV3(); + + expect(data.annotations.pfam).toEqual({ + kind: 'categorical', + // The encoded '|' - the grammar's suffix separator - is decoded back into a label. + values: ['PF00001 (7tm;1)', 'PF00002', 'PF00003 (a|b)', NA_VALUE], + colors: ['#F3C300', '#875692', '#F38400', NA_DEFAULT_COLOR], + shapes: ['circle', 'circle', 'circle', 'circle'], + }); + + expect(hitsByProtein(data, 'pfam')).toEqual({ + P1: { labels: [NA_VALUE], scores: [null], evidence: [] }, + // Two scores on one hit, one on the next: the score_count payload, not a 1:1 map. + P2: { labels: ['PF00001 (7tm;1)', 'PF00002'], scores: [[1e-10, 2.5], [0.5]], evidence: [] }, + P3: { labels: [NA_VALUE], scores: [null], evidence: [] }, + // P4 immediately follows the interior empty row: its score is what an off-by-one + // in the inserted-NA `hitEnd` would steal. + P4: { labels: ['PF00001 (7tm;1)'], scores: [[0.25]], evidence: [] }, + P5: { labels: ['PF00003 (a|b)'], scores: [[3]], evidence: [] }, + P6: { labels: [NA_VALUE], scores: [null], evidence: [] }, + }); + }); + + it('decodes a plain categorical column with no NA slot at all', async () => { + const { data } = await loadV3(); + + expect(data.annotations.kingdom).toEqual({ + kind: 'categorical', + values: ['Bacteria', 'Archaea', 'Eukaryota'], + colors: ['#F3C300', '#875692', '#F38400'], + shapes: ['circle', 'circle', 'circle'], + }); + // Every row has a kingdom, so no synthetic category may be appended. + expect(data.annotations.kingdom.values.some(isNAValue)).toBe(false); + expect(Array.from(data.annotation_data.kingdom as Int32Array)).toEqual([0, 1, 0, 2, 0, 1]); + }); + + it('folds every missing-value spelling in one dictionary into a single NA slot', async () => { + const { data } = await loadV3(); + + // The encoder is faithful: part 6 carries `none` (3 rows) and `NA` (1 row) as two + // ordinary labels, because collapsing them corrupts the Python side. The browser + // makes the display decision, and must land them in ONE bucket, not two. + expect(data.annotations.predicted_tm).toEqual({ + kind: 'categorical', + values: ['TM helix', NA_VALUE], + colors: ['#F3C300', NA_DEFAULT_COLOR], + shapes: ['circle', 'circle'], + }); + expect(data.annotations.predicted_tm.values.filter(isNAValue)).toHaveLength(1); + + const naIndex = data.annotations.predicted_tm.values.findIndex(isNAValue); + const codes = Array.from(data.annotation_data.predicted_tm as Int32Array); + expect(codes.filter((code) => code === naIndex)).toHaveLength(4); + // All three `none` rows plus the single `NA` row, and nothing else. + expect(codes).toEqual([naIndex, naIndex, 0, naIndex, naIndex, 0]); + }); + + it('reads numeric columns with their declared int/float type and null for blanks', async () => { + const { data } = await loadV3(); + + expect(data.annotations.length).toEqual({ + kind: 'numeric', + numericType: 'int', + values: [], + colors: [], + shapes: [], + }); + expect(data.annotations.hydrophobicity).toMatchObject({ + kind: 'numeric', + numericType: 'float', + }); + expect(data.numeric_annotation_data).toEqual({ + length: [120, null, 340, 0, -15, 1024], + hydrophobicity: [0.5, -1.25, null, 3, 0.001, 42], + }); + // A numeric column carries no categorical storage to bin by code. + expect(data.annotation_data.length).toBeUndefined(); + }); + + it('interleaves the wide axis columns into a 2D and a 3D projection', async () => { + const { data } = await loadV3(); + + expect(data.projections.map((projection) => projection.name)).toEqual(['pca2', 'umap3']); + const [pca2, umap3] = data.projections; + + expect(pca2.dimension).toBe(2); + expect(Array.from(pca2.data)).toEqual([0, 0, 1, 1, 2.5, -3.5, -4, 0.25, 5, 5, -1.5, 2]); + expect(pca2.metadata).toMatchObject({ components: 2, dimension: 2, dimensions: 2 }); + + expect(umap3.dimension).toBe(3); + expect(Array.from(umap3.data)).toEqual(Array.from({ length: 18 }, (_, index) => index / 4)); + expect(umap3.metadata).toMatchObject({ n_neighbors: 15, dimension: 3, dimensions: 3 }); + }); +}); + +describe('v3 -> v2 export round trip: CSR storage is interchangeable with nested', () => { + it('re-exports every categorical cell the encoder wrote, NA spellings excepted', async () => { + const { data } = await loadV3(); + const extraction = await extractRowsFromParquetBundle(createParquetBundle(data)); + + // The writer still stamps v2, so this proves the v2 cell grammar can be rebuilt from + // CSR storage: the reserved ';' and '|' inside labels go back out percent-encoded, + // scores rejoin with ',', evidence with '|'. + expect(extraction.formatVersion).toBe(2); + const cells = Object.fromEntries( + [...extraction.annotationsById].map(([id, row]) => [ + id, + Object.fromEntries(CATEGORICAL.map((key) => [key, row[key] ?? null])), + ]), + ); + const cathSemicolon = 'G3DSA:1.10.10.10 (Ribosomal Protein L15%3B Chain: K%3B domain 2)'; + expect(cells).toEqual({ + P1: { + cath: `${cathSemicolon}|50.2;G3DSA:6.20.10.10|60.5`, + go_bp: 'apoptotic process|IDA', + pfam: null, + kingdom: 'Bacteria', + // Documented non-identity: `none` is a MISSING_VALUE_TOKEN, so it was folded to + // `__NA__` on read and goes back out as NULL, not as the literal word. + predicted_tm: null, + }, + P2: { + cath: '6.20.10.10', + go_bp: null, + pfam: 'PF00001 (7tm%3B1)|1e-10,2.5;PF00002|0.5', + kingdom: 'Archaea', + predicted_tm: null, + }, + P3: { + cath: 'G3DSA:6.20.10.10|123456789', + go_bp: 'apoptotic process|IDA;protein folding|ECO:0000269', + pfam: null, + kingdom: 'Bacteria', + predicted_tm: 'TM helix', + }, + P4: { + cath: null, + go_bp: 'protein folding|IEA', + pfam: 'PF00001 (7tm%3B1)|0.25', + kingdom: 'Eukaryota', + predicted_tm: null, + }, + P5: { + cath: `${cathSemicolon}|1e-200`, + go_bp: null, + pfam: 'PF00003 (a%7Cb)|3', + kingdom: 'Bacteria', + // The other missing-value spelling in the same column, same treatment. + predicted_tm: null, + }, + P6: { + cath: '6.20.10.10', + go_bp: 'apoptotic process|EXP', + pfam: null, + kingdom: 'Archaea', + predicted_tm: 'TM helix', + }, + }); + // The in-memory sentinel is never written as a literal 6-char category. + for (const row of extraction.annotationsById.values()) { + expect(Object.values(row)).not.toContain(NA_VALUE); + } + }); + + it('reloads into the same legend, the same numerics and the same projections', async () => { + const { data: v3 } = await loadV3(); + const reloaded = await loadLegacy(createParquetBundle(v3)); + + expect(reloaded.protein_ids).toEqual(v3.protein_ids); + // Values, colours and shapes are what the legend and the palette are built from, so + // they have to survive a shape change that reorders nothing. + expect(reloaded.annotations).toEqual(v3.annotations); + expect(reloaded.numeric_annotation_data).toEqual(v3.numeric_annotation_data); + expect(reloaded.annotation_predicted).toBeUndefined(); + expect( + reloaded.projections.map(({ name, dimension, data }) => ({ + name, + dimension, + data: Array.from(data), + })), + ).toEqual( + v3.projections.map(({ name, dimension, data }) => ({ + name, + dimension, + data: Array.from(data), + })), + ); + }); + + it('reloads into nested storage that reads back identically through the accessors', async () => { + const { data: v3 } = await loadV3(); + const reloaded = await loadLegacy(createParquetBundle(v3)); + + // Precondition: the two datasets really are stored differently, otherwise the + // comparison below proves nothing about CSR at all. + expect(isCsrAnnotationData(v3.annotation_data.cath)).toBe(true); + expect(isCsrAnnotationData(reloaded.annotation_data.cath)).toBe(false); + expect(v3.annotation_data.kingdom).toBeInstanceOf(Int32Array); + expect(reloaded.annotation_data.kingdom).not.toBeInstanceOf(Int32Array); + + for (const key of CATEGORICAL) { + for (const [index, id] of v3.protein_ids.entries()) { + const from3 = hitsOf(v3, key, index); + const from2 = hitsOf(reloaded, key, index); + const where = `${key}/${id}`; + + expect(from2.labels, where).toEqual(from3.labels); + + // The ONE documented non-identity. An empty CSR row owns no hit slot, so the + // reader inserts a synthetic `__NA__` hit for it - and the flat score/evidence + // payloads are numbered by hit, so that inserted hit reports itself as `null`. + // Nested storage has no hit there at all and reports nothing. Asserted as the + // exact rows it applies to rather than by relaxing the comparison. + if (NA_ONLY_ROWS[key]?.includes(id)) { + expect(from3.labels, where).toEqual([NA_VALUE]); + const scored = key !== 'go_bp'; + expect(from3.scores, where).toEqual(scored ? [null] : []); + expect(from3.evidence, where).toEqual(scored ? [] : [null]); + expect(from2.scores, where).toEqual([]); + expect(from2.evidence, where).toEqual([]); + continue; + } + + expect(from2.scores, where).toEqual(from3.scores); + expect(from2.evidence, where).toEqual(from3.evidence); + } + } + }); +}); + +describe('v2 and v3 fixtures agree on the values they share', () => { + it('gives P1 and P2 the same ids, cath/go_bp hits and pca2 coordinates', async () => { + const { data: v3 } = await loadV3(); + const v2 = await loadLegacy(fixture('v2-sample.parquetbundle')); + + // The v3 fixture is a superset: 6 proteins to the v2 sample's 2, and 7 columns to + // its 2. Only the shared prefix is comparable. + expect(v2.protein_ids).toEqual(v3.protein_ids.slice(0, 2)); + + for (const [index, id] of v2.protein_ids.entries()) { + const cath3 = hitsOf(v3, 'cath', index); + expect(hitsOf(v2, 'cath', index), `cath/${id}`).toEqual(cath3); + + const go2 = hitsOf(v2, 'go_bp', index); + const go3 = hitsOf(v3, 'go_bp', index); + expect(go2.labels, `go_bp/${id}`).toEqual(go3.labels); + expect(go2.scores, `go_bp/${id}`).toEqual(go3.scores); + // P2's go_bp cell is empty in both fixtures, which is the same nested-vs-CSR + // non-identity the round trip above documents. + expect(go2.evidence, `go_bp/${id}`).toEqual(id === 'P2' ? [] : go3.evidence); + } + expect(hitsOf(v2, 'go_bp', 0).evidence).toEqual(['IDA']); + + const pca2v2 = v2.projections.find((projection) => projection.name === 'pca2')!; + const pca2v3 = v3.projections.find((projection) => projection.name === 'pca2')!; + expect(pca2v2.dimension).toBe(pca2v3.dimension); + expect(Array.from(pca2v2.data)).toEqual(Array.from(pca2v3.data.subarray(0, 4))); + }); +}); From bc01fb36d53453f1f0662061521381735ce60d85 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 03:40:10 +0200 Subject: [PATCH 29/31] docs(bundle): correct the v3 zero-copy, sourceType and score-spelling claims Also scope the 19%/21% figures and record the one-way version cross-check. --- apps/protspace/CLAUDE.md | 2 +- docs/guide/data-format.md | 76 +++++++++++++++++++++++++++++---------- 2 files changed, 58 insertions(+), 20 deletions(-) diff --git a/apps/protspace/CLAUDE.md b/apps/protspace/CLAUDE.md index efc85195..c359c3c5 100644 --- a/apps/protspace/CLAUDE.md +++ b/apps/protspace/CLAUDE.md @@ -248,7 +248,7 @@ HDF5 file (float16 embeddings) 3. `projections_data` — reduced coordinates per protein per projection 4. `settings` (optional) — annotation styles, pinned values, display config 5. `statistics` (optional) — tidy table of annotation-based validity (silhouette/DBI/CH per annotation, `space_kind ∈ {embedding, projection}`, `annotation` column) + auto-cluster ARI/NMI agreement (`stat_family=cluster_agreement`) (`protspace stats` / `prepare --stats`) -6. `payloads` (format v3 only, required) — label dictionaries and CSR code/score/evidence buffers for part 1 (`data/io/bundle_v3.py`) +6. `payloads` (format v3 only, required): label dictionaries and CSR code/score/evidence buffers for part 1 (`data/io/bundle_v3.py`) Every write from here emits **six** parts (format v3): `core(3) + settings + statistics + payloads`, with zero bytes in the settings or statistics slot when absent, because the browser reads the payloads positionally from `parts[5]`. `replace_settings_in_bundle` (`protspace style`) is the exception: it preserves the layout it was given, so a legacy bundle stays legacy. diff --git a/docs/guide/data-format.md b/docs/guide/data-format.md index 1a36be5f..1ba7b9df 100644 --- a/docs/guide/data-format.md +++ b/docs/guide/data-format.md @@ -18,8 +18,11 @@ metadata of its first part, under `protspace_format_version` (see | Columnar (format v3) | always 6 | `protspace prepare`, `protspace bundle`, `protspace transfer` | Both layouts carry the same data. v3 re-encodes the container, not the dataset: the Python API -decodes a v3 file back into exactly the three tables, with exactly the cell grammar, that a -legacy file stores directly. See [Format v3 Physical Schema](#format-v3-physical-schema). +decodes a v3 file back into the same three tables, in the same cell grammar, that a legacy file +stores directly. Cells come back in their canonical spelling rather than byte for byte, and the +deliberate differences are listed under +[What a v3 round trip does not preserve](#what-a-v3-round-trip-does-not-preserve). See +[Format v3 Physical Schema](#format-v3-physical-schema). ### Legacy layout (3 to 5 parts) @@ -323,6 +326,12 @@ When displayed in ProtSpace, the decoded names render as "Superfamily; old" and decode percent-encoded sequences. Existing v1 and v2 bundles therefore keep loading unchanged. - Python cross-checks the two signals: a six-part file whose first part does not say `3` is rejected rather than guessed at, and a three to five part file is always read as legacy. +- That cross-check runs in one direction only, which is a known limitation. Only a six-part file + has its version key consulted; a five-part file whose key says `3` is read as legacy without + complaint, and its part 1 comes back exactly as stored, `INT32` dictionary codes and + `__count` columns with no payloads part to resolve them against. Nothing ProtSpace writes + produces such a file, because every v3 write emits six parts, but a truncated or hand-assembled + one fails quietly instead of loudly. - The key versions the **container**, not the cell grammar. The grammar is still v2, which is why `BUNDLE_FORMAT_VERSION` on the Python side stays `2` and why the tables handed back from a v3 read are re-stamped `protspace_format_version=2`: what they contain is v2 cells. @@ -335,7 +344,10 @@ When displayed in ProtSpace, the decoded names render as "Superfamily; old" and Format v3 changes how the three logical tables are physically stored so that the browser can build its typed arrays without parsing a single annotation cell. On the 573K SwissProt dataset that takes bundle decoding from about 6.5 s and 2.1 GB of heap to about 0.4 s and under 50 MB, -in a file about 19% smaller than the v2 encoding of the same data. +in a file about 19% smaller than the v2 encoding of the same data. The module docstring in +`bundle_v3.py` quotes 21% for the same bundle. The two numbers measure different builds, not the +same one twice: 21% is what settled the counts-versus-offsets choice described below, measured +while scores were still float32, and 19% is the same comparison after scores widened to float64. Nothing above the container boundary changes. `read_tables()` decodes a v3 file back into the v2-shaped tables, so `protspace serve`, `protspace style`, `protspace transfer` and every script @@ -401,7 +413,7 @@ label's bytes. The reader prefix-sums them into offsets in a single pass. That is a size decision. Offsets are near incompressible (snappy manages about 0.4% of them on the 573K SwissProt bundle) while their first differences, which is what the counts are, compress -about 8x. On that bundle the difference is roughly 15 MB, about 10.0 MB of it in part 1 and +about 8x. On that bundle the difference is roughly 16 MB, about 10.0 MB of it in part 1 and about 5.7 MB in part 6. ### Required, PLAIN, one row group @@ -409,10 +421,18 @@ about 5.7 MB in part 6. Every column of parts 1, 3 and 6 is written non-nullable, PLAIN encoded, with dictionary encoding disabled, in a single row group, snappy compressed. -That is load bearing, not stylistic. The browser's Parquet reader hands back a zero-copy typed -array only for a REQUIRED flat PLAIN column. A column written nullable or dictionary-encoded +For parts 1 and 3 that is load bearing, not stylistic. The browser's Parquet reader hands back a +zero-copy typed array only for a REQUIRED flat PLAIN column, and each such chunk then lands in +its preallocated column with a single `set`. A column written nullable or dictionary-encoded still decodes to the right values, but it arrives as a plain JavaScript array about 4x slower, -and the reader logs one warning naming it. Such a column is a writer bug, not a variant. +and the reader logs a warning naming it. Parts 1 and 3 are read by separate passes that carry one +warning each, so a single read logs at most two. Such a column is a writer bug, not a variant. + +Part 6 is read differently and is not zero-copy at all. It is a handful of large blobs rather +than hundreds of per-row columns, so the reader loads the whole part at once and copies each +payload out of the decoded page before wrapping it as a typed array. At that granularity the copy +costs almost nothing, and it is what guarantees alignment: the payload arrives as a view at an +arbitrary byte offset, which a `Float64Array` cannot wrap. ### The manifest @@ -447,9 +467,21 @@ code column declared numeric, for instance), a duplicate projection name, or a d not 2 or 3 all throw rather than being repaired. `sourceType` is Python-private and the browser ignores it. It records the Arrow type the column -had before encoding, so the decoder can restore it instead of handing back a string column, and -it is `"?"` for a type that cannot be parsed back from its alias, such as a dictionary, list or -decimal column. The decoder falls back to the per-kind default for those. +had before encoding, so the decoder can restore it instead of handing back a string column. + +Only a numeric source type is ever restored, and only on a numeric column. The decoder declines +everything else and falls back to the per-kind default: + +- a type whose alias cannot be parsed back at all, such as a dictionary, list or decimal column, + is recorded as `"?"` +- a type whose alias parses but is not an integer or a float is recorded verbatim and declined + anyway. A `bool` column records `sourceType` `"bool"` and a timestamp column records + `"timestamp[s]"`; neither is restored, and both come back as v2 string cells +- `"string"` and `"large_string"` are declined by design, because the v2 spelling of the column is + what the encoder consumed, so rendering it back is the restoration + +`sourceType` is also consulted only when decoding a `numeric` column, so on a `categorical` or +`multi` column it is recorded but completely inert. ### Scores are float64 @@ -481,15 +513,21 @@ A v3 file stores what the reader would have parsed out of the v2 cells, not the so decoding a v3 bundle returns the canonical spelling of each cell rather than the original bytes. The differences are deliberate: -| Written | Read back | Why | -| ------------------------------------- | --------------- | ------------------------------------------------------------------------------------------------- | -| `PF00001\|62.0` | `PF00001\|62` | scores are re-spelled the way JavaScript prints them, and JavaScript never prints a trailing `.0` | -| `PF00001\|0.5700` | `PF00001\|0.57` | same rule: the shortest spelling that reads back as the same double | -| ` A \|IDA` | `A\|IDA` | cells and hits are whitespace-trimmed | -| `A;;B` | `A;B` | empty hits are dropped | -| `%3b` | `%3B` | labels are re-encoded canonically, in uppercase hex | -| a missing cell | `""` | null and blank both mean missing | -| `100.0` where every value is integral | `100` | the canonical v2 spelling of an integral value | +| Written | Read back | Why | +| ------------------------------------- | ------------------ | ------------------------------------------------------------------------------------------------------------------------------- | +| `PF00001\|62.0` | `PF00001\|62` | scores are re-spelled by Python's float repr, the shortest spelling that reads back as the same double, minus the trailing `.0` | +| `PF00001\|0.5700` | `PF00001\|0.57` | same rule | +| `PF00001\|2.3e-5` | `PF00001\|2.3e-05` | same rule; Python's repr is not JavaScript's, and it pads a single-digit exponent to two digits | +| `PF00001\|1e16` | `PF00001\|1e+16` | same rule; Python switches to exponential notation at 1e16, JavaScript not until 1e21 | +| ` A \|IDA` | `A\|IDA` | cells and hits are whitespace-trimmed | +| `A;;B` | `A;B` | empty hits are dropped | +| `%3b` | `%3B` | labels are re-encoded canonically, in uppercase hex | +| a missing cell | `""` | null and blank both mean missing | +| `100.0` where every value is integral | `100` | the canonical v2 spelling of an integral value | + +The `2.3e-5` and `1e16` rows land squarely in the E-value range the float64 scores exist for, so +they are worth knowing about, but they are cosmetic rather than corrupting: both spellings +re-parse to the same double, and a second round trip re-emits the same text. A cell spelled `none`, `NA` or `null` is an ordinary label and comes back unchanged. Projection coordinates come back as float32 with `z` null for a 2D projection, and the identifier column From 597afbbf487bfded62b199eb99de750a7dabf243 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 03:58:24 +0200 Subject: [PATCH 30/31] fix(bundle): guard the v3 write path's unstamped-table and corrupt-part hazards Cover the annotations re-stamp, warn on an unstamped encode, drop empty-string legend buckets, and correct three stale v3 comments. --- .../protspace/src/protspace/data/io/bundle.py | 29 +++++++- .../src/protspace/data/io/bundle_v3.py | 53 +++++++++++-- .../src/protspace/utils/arrow_reader.py | 12 ++- .../tests/test_bundle_v3_container.py | 74 +++++++++++++++++++ apps/protspace/tests/test_display_decode.py | 38 ++++++++++ apps/protspace/tests/test_transfer_cli.py | 8 +- 6 files changed, 200 insertions(+), 14 deletions(-) diff --git a/apps/protspace/src/protspace/data/io/bundle.py b/apps/protspace/src/protspace/data/io/bundle.py index 64ef8ab0..57601c3d 100644 --- a/apps/protspace/src/protspace/data/io/bundle.py +++ b/apps/protspace/src/protspace/data/io/bundle.py @@ -49,8 +49,19 @@ def _part_container_version(part: bytes) -> int: - """The ``protspace_format_version`` in a part's parquet footer (1 if absent).""" - metadata = pq.read_metadata(io.BytesIO(part)).metadata or {} + """The ``protspace_format_version`` in a part's parquet footer (1 if absent). + + Every six-part read parses part 1's footer, including the settings-only + :func:`read_settings_from_bundle`, so a corrupt part 1 has to fail as a + bundle error and not as a raw ``ArrowInvalid`` traceback out of + ``protspace style --dump-settings``. + """ + try: + metadata = pq.read_metadata(io.BytesIO(part)).metadata or {} + except pa.ArrowInvalid as exc: + raise ValueError( + f"parquetbundle part 1 is not readable as parquet: {exc}" + ) from exc try: return int(metadata.get(FORMAT_VERSION_KEY, b"1")) except (TypeError, ValueError): @@ -280,6 +291,20 @@ def write_bundle( The tables come in v2-shaped (all-string annotation cells, long-format projections) and go out as a six-part v3 container. + **Precondition: ``tables[0]`` must carry the format-version stamp** unless it + really is v1. An unstamped table reads back as v1 (:func:`read_format_version` + defaults it), so :func:`~protspace.data.io.bundle_v3.encode_v3` migrates it -- + and migrating an already-v2 table double-escapes every reserved character + (``%3B`` becomes ``%253B``), unrecoverably, because ``decode_field`` is not + its own inverse. pyarrow drops schema metadata on ``rename_columns``, + ``concat_tables`` and friends, so a caller that rebuilds the table must + re-apply :func:`~protspace.data.annotations.encoding.stamp_format_version` + afterwards, as ``cli/bundle.py`` does. ``encode_v3`` warns instead of + refusing, and this function cannot stamp for its callers the way + :func:`replace_annotations_in_bundle` does: it is also the path a genuine + legacy bundle is upgraded through, and there the unstamped table really is + v1. + Args: tables: List of 3 Arrow tables (annotations, projections_metadata, projections_data). diff --git a/apps/protspace/src/protspace/data/io/bundle_v3.py b/apps/protspace/src/protspace/data/io/bundle_v3.py index abc43f01..b3412ed7 100644 --- a/apps/protspace/src/protspace/data/io/bundle_v3.py +++ b/apps/protspace/src/protspace/data/io/bundle_v3.py @@ -30,6 +30,7 @@ import io import json +import logging from typing import Any import numpy as np @@ -46,6 +47,8 @@ stamp_format_version, ) +logger = logging.getLogger(__name__) + CONTAINER_VERSION = 3 MANIFEST_KEY = b"protspace_v3_manifest" @@ -478,8 +481,25 @@ def encode_v3( projections_metadata: pa.Table, projections_data: pa.Table, ) -> tuple[bytes, bytes, bytes, bytes]: - """Encode the v2-shaped pipeline tables as v3 parts 1, 2, 3 and 6.""" + """Encode the v2-shaped pipeline tables as v3 parts 1, 2, 3 and 6. + + ``annotations`` must carry the format-version stamp unless it really is v1: + an unstamped v2 table is indistinguishable from a v1 one here and is + migrated a second time, double-escaping every reserved character. See the + precondition on :func:`~protspace.data.io.bundle.write_bundle`. + """ if read_format_version(annotations) == 1: + # Loud, because the alternative failure is silent and unrecoverable: an + # already-v2 table that lost its stamp (pyarrow drops schema metadata on + # rename_columns/concat) is migrated twice and every ``%3B`` becomes + # ``%253B``. Refusing instead is not an option -- a genuine legacy + # bundle read back is unstamped too, and that upgrade path is the point. + logger.warning( + "annotations table reads as format v1 (no stamp, or stamped 1); " + "migrating its cell grammar to v2. If it was already v2, re-apply " + "stamp_format_version() before writing -- migrating twice escapes " + "every reserved character a second time." + ) annotations = migrate_legacy_annotation_table(annotations) id_column = next( @@ -718,10 +738,16 @@ def _decode_multi( # InterPro score) underflows float32 to ``0`` and ``1e40`` overflows to # ``inf``, which is not even valid v2, so a second round trip would # re-classify the hit. numpy's float64 repr is the shortest spelling that - # reads back as the same double, which is what ``String(number)`` gives - # the browser -- bar - # the trailing ``.0`` JavaScript never prints (``[1].join(',')`` is - # ``"1"``), so ``Array.prototype.join`` and this agree on every score. + # reads back as the same double, but it is *not* character-identical to + # ``String(number)``: Python pads the exponent to two digits (``1e-07`` + # where JavaScript prints ``1e-7``) and switches to exponential notation + # at different magnitudes -- below 1e-4 against JavaScript's 1e-6 + # (``2.3e-05`` against ``0.000023``) and from 1e16 against JavaScript's + # 1e21. Both spellings parse back to the same double, so the difference + # is cosmetic in a score suffix and never changes a comparison; only the + # trailing ``.0`` is normalised away here, because JavaScript never + # prints it (``[1].join(',')`` is ``"1"``) and it would otherwise shift + # the label text. text = pc.replace_substring_regex( pa.array(values.astype(str), type=pa.string()), r"\.0$", "" ) @@ -791,8 +817,14 @@ def decode_v3(parts: list[bytes]) -> tuple[pa.Table, pa.Table, pa.Table]: The round trip is not byte-exact, and deliberately so -- v3 stores what the browser's v2 reader would have parsed out of the cells, not the cells: - * hits and cells are whitespace-trimmed, and empty or missing-valued hits are - dropped (``"A;;B"`` comes back ``"A;B"``, ``" A |IDA"`` as ``"A|IDA"``); + * an unstamped (v1-reading) input table is migrated to the v2 cell grammar + first, which is what every shipped legacy bundle gets: ten of the eleven + datasets under ``apps/web/public/data/`` carry no stamp, so regenerating + one runs :func:`migrate_legacy_annotation_table` and its reserved + characters come back percent-encoded (display-neutral, and a fix -- but it + is a difference, and it is the one a regeneration actually hits); + * hits and cells are whitespace-trimmed, and *blank* hits are dropped + (``"A;;B"`` comes back ``"A;B"``, ``" A |IDA"`` as ``"A|IDA"``); * a missing cell -- null or blank -- comes back as ``""`` (a cell spelled ``none``/``NA``/``null`` is an ordinary label and comes back unchanged); * labels are re-encoded canonically, so ``%3b`` comes back as ``%3B``; @@ -801,6 +833,13 @@ def decode_v3(parts: list[bytes]) -> tuple[pa.Table, pa.Table, pa.Table]: * a numeric column comes back in its ``sourceType`` when that is restorable and otherwise as its canonical v2 spelling, so an all-integral column spells ``100``, never ``100.0``; + * a non-finite value in an Arrow-numeric column is **lost**: ``±inf`` and + ``NaN`` both encode as null and come back null (or ``""``). Unreachable + from ``prepare`` -- nothing upstream emits one -- but ``protspace bundle + -a `` will happily hand one over. Neither is expressible in + v2 anyway: the cell grammar's number rule rejects ``Infinity`` and the + browser drops non-finite values on read, so preserving them would produce + a bundle no reader agrees on; * projection coordinates come back float32 (``z`` null for a 2D projection) and a protein absent from a projection comes back at the origin; * the identifier column comes back first, wherever it sat before. diff --git a/apps/protspace/src/protspace/utils/arrow_reader.py b/apps/protspace/src/protspace/utils/arrow_reader.py index 32f88e67..b8d8f537 100644 --- a/apps/protspace/src/protspace/utils/arrow_reader.py +++ b/apps/protspace/src/protspace/utils/arrow_reader.py @@ -295,11 +295,19 @@ def get_marker_shape(self, annotation: str) -> dict[str, str]: ) def get_unique_annotation_values(self, annotation: str) -> list[Any]: - """Get a list of unique values for a given annotation.""" + """Get a list of unique values for a given annotation. + + Absent values are skipped, and ``""`` counts as absent: a v3 container + spells a missing categorical cell as the empty string (a documented + ``decode_v3`` non-identity), so keeping it would add an empty bucket to + every column that has any missing value -- 427 of them on + ``venom_eat_stats``'s ``ec__pred_value`` alone, which nothing filters out + of the annotation list. + """ unique_values = set() for protein_data in self.data.get("protein_data", {}).values(): value = protein_data.get("annotations", {}).get(annotation) - if value is not None: + if value is not None and value != "": unique_values.add(value) return list(unique_values) diff --git a/apps/protspace/tests/test_bundle_v3_container.py b/apps/protspace/tests/test_bundle_v3_container.py index ed4f7ddb..bf32967c 100644 --- a/apps/protspace/tests/test_bundle_v3_container.py +++ b/apps/protspace/tests/test_bundle_v3_container.py @@ -16,6 +16,7 @@ """ import io +import logging from pathlib import Path import pyarrow as pa @@ -32,6 +33,7 @@ _write_parts, extract_bundle_to_dir, read_bundle, + read_settings_from_bundle, read_tables, replace_annotations_in_bundle, replace_settings_in_bundle, @@ -292,6 +294,78 @@ def test_replace_annotations_re_encodes_the_payloads(tmp_path): assert b"Bacteria" not in payload_blob # no stale dictionary left behind +def test_replace_annotations_keeps_encoded_cells_from_an_unstamped_table(tmp_path): + """The re-stamp in ``replace_annotations_in_bundle`` is what stops a second + migration. ``transfer`` and the prediction overlay rebuild the table with + ``rename_columns``/``concat_tables``, which drop schema metadata, so what + arrives here is v2 cells that *read* as v1 -- and migrating them again turns + ``%3B`` into ``%253B``, unrecoverably (``decode_field`` is not its own + inverse). Drop the ``stamp_format_version`` line at the chokepoint and this + fails on both cells; nothing else in the suite does.""" + src = tmp_path / "b.parquetbundle" + out = tmp_path / "out.parquetbundle" + write_bundle(pipeline_tables(), src) + + cells = ["G3DSA:1.1 (Ribosomal L15%3B Chain: K)", "PF1 (a%7Cb)|0.5", "plain"] + unstamped = pa.table({"protein_id": ["p0", "p1", "p2"], "cath": cells}) + assert FORMAT_VERSION_KEY not in (unstamped.schema.metadata or {}) + + replace_annotations_in_bundle(src, out, unstamped) + + assert read_tables(out)[0].column("cath").to_pylist() == cells + + +def test_write_bundle_warns_when_the_annotations_table_is_unstamped(tmp_path, caplog): + """``write_bundle`` has no chokepoint re-stamp and cannot have one: it is + also the path a genuine legacy bundle is upgraded through, and an unstamped + v1 table is indistinguishable from a v2 one that lost its metadata. So the + migration is at least loud -- the docstring states the precondition and the + encoder says out loud which of the two it decided it got. The second half + of the assertion is the cost of getting it wrong.""" + metadata, data = projection_tables(2, (2,)) + unstamped = pa.table({"protein_id": ["p0", "p1"], "cath": ["ACC (a%3Bb)", "x"]}) + path = tmp_path / "b.parquetbundle" + + with caplog.at_level(logging.WARNING, logger="protspace.data.io.bundle_v3"): + write_bundle([unstamped, metadata, data], path) + + assert "format v1" in caplog.text + # Migrated a second time, because it looked like v1: the caller was warned. + assert read_tables(path)[0].column("cath").to_pylist() == ["ACC (a%253Bb)", "x"] + + caplog.clear() + with caplog.at_level(logging.WARNING, logger="protspace.data.io.bundle_v3"): + write_bundle([stamp_format_version(unstamped), metadata, data], path) + + assert caplog.text == "" + assert read_tables(path)[0].column("cath").to_pylist() == ["ACC (a%3Bb)", "x"] + + +def test_write_bundle_rejects_a_fourth_table(tmp_path): + """Untested until now: removing the guard passes the whole suite. Without + it the tuple unpack below still raises, but as a bare "too many values to + unpack" naming neither the function nor what it wanted.""" + tables = pipeline_tables() + with pytest.raises(ValueError, match="expects 3 core tables"): + write_bundle([*tables, tables[0]], tmp_path / "b.parquetbundle") + + +def test_corrupt_first_part_of_a_six_part_bundle_is_a_bundle_error(tmp_path): + """Every six-part read parses part 1's footer, ``read_settings_from_bundle`` + included -- so a corrupt part 1 must not surface as a raw ``ArrowInvalid`` + traceback out of ``protspace style --dump-settings``, which never touched + part 1 before.""" + path = tmp_path / "b.parquetbundle" + write_bundle(pipeline_tables(), path, settings={"k": 1}) + parts = parts_of(path) + path.write_bytes(PARQUET_BUNDLE_DELIMITER.join([b"not parquet", *parts[1:]])) + + with pytest.raises(ValueError, match="part 1 is not readable as parquet"): + read_settings_from_bundle(path) + with pytest.raises(ValueError, match="part 1 is not readable as parquet"): + read_tables(path) + + def test_replace_annotations_upgrades_a_legacy_bundle(tmp_path): """A rewrite is a write, and every write emits v3.""" src = tmp_path / "legacy.parquetbundle" diff --git a/apps/protspace/tests/test_display_decode.py b/apps/protspace/tests/test_display_decode.py index 06576d65..6c18770a 100644 --- a/apps/protspace/tests/test_display_decode.py +++ b/apps/protspace/tests/test_display_decode.py @@ -212,3 +212,41 @@ def test_style_bundle_roundtrips_for_encoded_name(tmp_path): # keyed by the decoded display name; #123456 is normalized to its rgba form assert display_name in colors assert "18, 52, 86" in colors[display_name] + + +def test_unique_annotation_values_skip_empty_strings(tmp_path): + """A v3 container spells a missing categorical cell as ``""``, so the Dash + reader sees a value where a null used to be. Nothing downstream filters it + out, so an unskipped ``""`` becomes an extra (blank) entry in the annotation + value list -- 427 of them on the shipped ``venom_eat_stats``'s + ``ec__pred_value``, which gains a fifth "value" it never had.""" + ann = stamp_format_version( + pa.table( + { + "protein_id": ["P1", "P2", "P3"], + "ec__pred_value": ["1.1.1.1", None, "2.2.2.2"], + } + ) + ) + meta = pa.table( + {"projection_name": ["pca2"], "dimensions": [2], "info_json": ["{}"]} + ) + data = pa.table( + { + "projection_name": ["pca2"] * 3, + "identifier": ["P1", "P2", "P3"], + "x": [0.0, 1.0, 2.0], + "y": [0.0, 1.0, 2.0], + "z": [None, None, None], + } + ) + path = tmp_path / "d.parquetbundle" + write_bundle([ann, meta, data], path) + + reader = ArrowReader(Path(extract_bundle_to_dir(path))) + # The null really did come back as "" -- the reason this test exists. + assert reader.get_protein_annotations("P2")["ec__pred_value"] == "" + assert sorted(reader.get_unique_annotation_values("ec__pred_value")) == [ + "1.1.1.1", + "2.2.2.2", + ] diff --git a/apps/protspace/tests/test_transfer_cli.py b/apps/protspace/tests/test_transfer_cli.py index 4950ddce..1af76972 100644 --- a/apps/protspace/tests/test_transfer_cli.py +++ b/apps/protspace/tests/test_transfer_cli.py @@ -575,7 +575,9 @@ def test_cli_transfer_without_rules_fills_missing_values(tmp_path): assert rows["TRINITY_1"]["protein_category__pred_value"] == "neurotoxin" # A reference protein gets no prediction. The overlay writes null; a v3 # container stores "absent" as a -1 dictionary code and spells it back as "" - # (a documented decode_v3 non-identity). Both readers of this column treat - # null and "" identically -- the browser's readCategoricalStorageValue and - # Python's to_display_value -- so the distinction is not observable. + # (a documented decode_v3 non-identity). The browser cannot tell the two + # apart (readCategoricalStorageValue folds both into missing) but Dash can: + # ArrowReader hands the raw cell out, so an "" is a value where a null was + # not. get_unique_annotation_values skips "" for exactly that reason -- see + # test_unique_annotation_values_skip_empty_strings. assert rows["P00001"]["protein_category__pred_value"] == "" From eb237b14ce4dbce779a954f0c9c519616a413e13 Mon Sep 17 00:00:00 2001 From: peymanvahidi Date: Sun, 6 Sep 2026 04:08:54 +0200 Subject: [PATCH 31/31] test(bundle): make the v3 golden fixture able to fail Encoded-part drift guard plus fixture cells that separate dictionary order, both payload families, non-ASCII labels and the EAT overlay. --- apps/protspace/scripts/generate_v3_fixture.py | 77 ++++- .../protspace/tests/test_bundle_v3_fixture.py | 225 ++++++++++++-- .../__fixtures__/v3-sample.parquetbundle | Bin 12809 -> 13532 bytes .../data-loader/utils/v3-equivalence.test.ts | 289 +++++++++++++++--- 4 files changed, 510 insertions(+), 81 deletions(-) diff --git a/apps/protspace/scripts/generate_v3_fixture.py b/apps/protspace/scripts/generate_v3_fixture.py index c3885865..0352987f 100644 --- a/apps/protspace/scripts/generate_v3_fixture.py +++ b/apps/protspace/scripts/generate_v3_fixture.py @@ -7,28 +7,43 @@ ``v2-sample.parquetbundle``: proteins ``P1``/``P2`` carry byte-identical ``cath`` and ``go_bp`` cells (including the percent-encoded ``;`` inside a CATH name and the ``|IDA`` evidence suffix), so every assertion the v2 golden test -makes still holds, and four more proteins plus six more columns cover the v3 +makes still holds, and four more proteins plus seven more columns cover the v3 paths a two-row two-column table cannot reach: -* a plain single-valued categorical (``kingdom``) whose frequency order differs - from its first-occurrence order; +* a single-valued categorical (``kingdom``, one cell blank) whose descending-frequency + dictionary order (``Bacteria``, ``Archaea``, ``Eukaryota``) really does differ + from its first-occurrence order (``Archaea``, ``Bacteria``, ``Eukaryota``), so + deleting the encoder's frequency sort changes the committed bytes. Dictionary + order is legend order and therefore colour assignment, and nothing else in the + fixture can tell the two orderings apart; * multi-valued columns with scores (``cath``, ``pfam``) and with evidence codes (``go_bp``, both the ``IDA`` and the ``ECO:0000269`` spellings); -* ``pfam`` also has two scores on one hit, and zero hits at the first, an - interior and the last row, so the reader's CSR prefix-sum and the synthetic - ```` insertion are exercised at every boundary; +* ``pfam`` is the one column that carries both payload families at once, scores + *and* evidence; it also has two scores on one hit, zero + hits at the first, an interior and the last row, a hit whose label is the + missing-value spelling ``none`` (the browser's only chance to run + ``dropFoldedHits``), a non-ASCII label that forces the browser's dictionary + reader off its pure-ASCII fast path onto per-label byte slicing, a score + written ``62.0`` and one written ``2.3e-5``; * scores that only survive in float64: ``1e-200`` flushes to zero in float32 and ``123456789`` re-spells as ``1.2345679e+08``; * a numeric int column (``length``) and a numeric float one (``hydrophobicity``), each with a blank cell; +* ``reviewed``, the one categorical with no gap at all, so the synthetic + ```` legend row must *not* be appended to it (its dictionary order also + disagrees with its first-occurrence order); * ``predicted_tm``, whose labels are the literal missing-value spellings ``none`` and ``NA`` — Python keeps them (v3 is a container encoding), the browser folds them into ```` at read time; * a 2D (``pca2``, P1/P2 at the v2 fixture's coordinates) and a 3D (``umap3``) - projection; + projection, with P6 absent from ``umap3`` so the 0.0-at-origin fill for a + protein missing from a projection is exercised; * the EAT companion trio on ``kingdom`` (``__pred_value`` string, ``__pred_confidence`` float32, ``__pred_source`` string), null for the - proteins with no prediction. + proteins with no prediction. Only P4's prediction survives the overlay's + "a curated value wins" rule, because P4 is the one protein whose curated + ``kingdom`` is blank, and its label ``Viruses`` appears nowhere in the curated + column, so the prediction-only legend entry is exercised too. Settings and statistics are deliberately absent, so the container's two zero-byte slots keep the payloads part at position six. @@ -46,6 +61,7 @@ import numpy as np import pandas as pd import pyarrow as pa +import pyarrow.compute as pc from protspace.data.annotations.encoding import encode_field, stamp_format_version from protspace.data.io.bundle import write_bundle @@ -71,6 +87,11 @@ _CATH_NAME = encode_field("Ribosomal Protein L15; Chain: K; domain 2") _CATH_1 = f"G3DSA:1.10.10.10 ({_CATH_NAME})" +# A label outside ASCII: its UTF-8 byte length is not its JavaScript string length, +# so the browser has to slice this dictionary per label by byte range instead of +# decoding the whole blob once. Every other dictionary here is pure ASCII. +PFAM_NON_ASCII_LABEL = "PF00004 (\u03b2-lactamase, N\u00e9buline)" + ANNOTATION_CELLS: dict[str, list[str]] = { # P1/P2 are the v2 fixture's cells, verbatim. "cath": [ @@ -91,17 +112,30 @@ ], # Zero hits first, interior and last; two scores on one hit; a label # carrying the encoded '|' the grammar reserves as the suffix separator. + # P4 mixes a scored hit, an EVIDENCE hit -- so this one column carries both + # payload families, which no other column crosses -- and a hit spelled + # `none`, the only folded missing-value label inside a multi column. + # P5 pins two score spellings the decode is documented to change: `62.0` + # comes back `62`, and `2.3e-5` comes back `2.3e-05` from Python and + # `0.000023` from the browser's exporter. "pfam": [ "", "PF00001 (7tm%3B1)|1e-10,2.5;PF00002|0.5", "", - "PF00001 (7tm%3B1)|0.25", - "PF00003 (a%7Cb)|3", + f"PF00001 (7tm%3B1)|0.25;{PFAM_NON_ASCII_LABEL}|IDA;none", + "PF00003 (a%7Cb)|3;PF00002|62.0;PF00001 (7tm%3B1)|2.3e-5", "", ], - # Frequency order (Bacteria, Archaea, Eukaryota) differs from first - # occurrence only at the tail, which is what pins the tie-break rule. - "kingdom": ["Bacteria", "Archaea", "Bacteria", "Eukaryota", "Bacteria", "Archaea"], + # Descending frequency (Bacteria 3, then Archaea 1 and Eukaryota 1) puts + # Bacteria first; first occurrence puts Archaea first. The two orderings + # DISAGREE, which is the only thing that can catch a lost frequency sort. + # The blank cell is the one curated gap the EAT overlay is allowed to fill. + "kingdom": ["Archaea", "Bacteria", "Bacteria", "", "Bacteria", "Eukaryota"], + # The one categorical with no gap anywhere: no blank cell and no + # missing-value spelling, so the browser must NOT append a synthetic + # legend row to it. Its frequency order (True, False) also disagrees with + # its first-occurrence order (False, True). + "reviewed": ["False", "True", "True", "False", "True", "True"], # Literal missing-value spellings kept as labels: a display decision the # browser makes, not a container one. "predicted_tm": ["none", "none", "TM helix", "none", "NA", "TM helix"], @@ -112,7 +146,7 @@ # (query_id, label, reliability, distance, source_id) for the EAT overlay. PREDICTIONS = [ ("P2", "Bacteria", 0.875, 0.12, "Q9XYZ1"), - ("P4", "Archaea", 0.5, 0.44, "P0A7B8"), + ("P4", "Viruses", 0.5, 0.44, "P0A7B8"), ("P5", "Bacteria", 0.25, 0.91, "A0A123"), ] @@ -171,10 +205,23 @@ def source_tables() -> list[pa.Table]: # append_column/drop_columns are not guaranteed to carry schema metadata. annotations = stamp_format_version(annotations) + coordinates = processor._create_projections_data_table(PROJECTIONS, PROTEIN_IDS) + # P6 has no umap3 row at all: the encoder fills 0.0 for a protein missing + # from a projection and the browser leaves its zero-initialised slot alone, + # so both put P6 at the origin. Pinned as a quirk, not endorsed. + coordinates = coordinates.filter( + pc.invert( + pc.and_( + pc.equal(coordinates.column("projection_name"), "umap3"), + pc.equal(coordinates.column("identifier"), "P6"), + ) + ) + ) + return [ annotations, processor._create_projections_metadata_table(PROJECTIONS), - processor._create_projections_data_table(PROJECTIONS, PROTEIN_IDS), + coordinates, ] diff --git a/apps/protspace/tests/test_bundle_v3_fixture.py b/apps/protspace/tests/test_bundle_v3_fixture.py index 14fc5cf1..162c3e3b 100644 --- a/apps/protspace/tests/test_bundle_v3_fixture.py +++ b/apps/protspace/tests/test_bundle_v3_fixture.py @@ -7,6 +7,12 @@ exact cells the browser side asserts on -- a fixture nobody reads back in Python is a fixture that silently rots. +The drift guard compares the **encoded** parts, not only the decoded tables. A +decoded comparison is blind to every encoder change the decoder symmetrically +undoes, and dictionary order is exactly that: drop the encoder's +descending-frequency sort and ``read_tables`` still hands back the same cells, +while the browser's legend and every colour in it silently reorder. + Regenerate with ``uv run python scripts/generate_v3_fixture.py``. """ @@ -16,6 +22,8 @@ import sys from pathlib import Path +import numpy as np +import pyarrow as pa import pyarrow.parquet as pq import pytest @@ -24,7 +32,7 @@ read_settings_from_bundle, read_tables, ) -from protspace.data.io.bundle_v3 import MANIFEST_KEY +from protspace.data.io.bundle_v3 import MANIFEST_KEY, decode_v3, encode_v3 FIXTURE = ( Path(__file__).resolve().parents[3] @@ -38,6 +46,12 @@ / "v3-sample.parquetbundle" ) +#: ``packages/core``. The only legitimate reason for the fixture to be absent is +#: that the browser workspace is not checked out at all (a Python-only source +#: distribution); a *deleted* fixture inside a present workspace is a deleted +#: contract and must fail, not skip. +_BROWSER_PACKAGE = FIXTURE.parents[5] + _SPEC = importlib.util.spec_from_file_location( "generate_v3_fixture", Path(__file__).parent.parent / "scripts" / "generate_v3_fixture.py", @@ -47,20 +61,75 @@ _SPEC.loader.exec_module(generator) pytestmark = pytest.mark.skipif( - not FIXTURE.exists(), reason="browser fixtures not checked out" + not _BROWSER_PACKAGE.is_dir(), reason="packages/core is not checked out" ) +#: The one label in the fixture whose UTF-8 byte length is not its character +#: count, which is what forces the browser off its pure-ASCII dictionary path. +NON_ASCII_LABEL = generator.PFAM_NON_ASCII_LABEL + + +def _table(part: bytes) -> pa.Table: + return pq.read_table(io.BytesIO(part)) + + +def _manifest(part: bytes) -> dict: + return json.loads(pq.read_metadata(io.BytesIO(part)).metadata[MANIFEST_KEY]) + + +def _payload_map(part: bytes) -> dict[str, bytes]: + """Part 6 as the ``name -> bytes`` map the browser builds from it.""" + table = _table(part) + return dict( + zip( + table.column("name").to_pylist(), + table.column("data").to_pylist(), + strict=True, + ) + ) + + +def _i32(payloads: dict[str, bytes], name: str) -> list[int]: + return np.frombuffer(payloads[name], " list[str]: + """Decode one dictionary payload the way the browser does: by byte length.""" + blob = payloads[f"dict:{name}"] + ends = np.cumsum(_i32(payloads, f"dict:{name}:len")) + return [ + blob[end - length : end].decode() + for length, end in zip(_i32(payloads, f"dict:{name}:len"), ends, strict=True) + ] + @pytest.fixture(scope="module") def parts() -> list[bytes]: return FIXTURE.read_bytes().split(PARQUET_BUNDLE_DELIMITER) +@pytest.fixture(scope="module") +def payloads(parts) -> dict[str, bytes]: + return _payload_map(parts[5]) + + @pytest.fixture(scope="module") def tables(): return read_tables(FIXTURE) +def test_the_fixture_is_committed(): + """A missing fixture is a deleted contract, not a reason to skip. + + Every other test here depends on the file, so this one names the failure + instead of leaving eight identical ``FileNotFoundError`` tracebacks. + """ + assert FIXTURE.exists(), ( + f"{FIXTURE} is missing while its package is checked out; regenerate it " + "with `uv run python scripts/generate_v3_fixture.py`" + ) + + def test_fixture_is_a_six_part_v3_container(parts): assert len(parts) == 6 # No settings, no statistics -- the two zero-byte slots are what keeps the @@ -70,11 +139,11 @@ def test_fixture_is_a_six_part_v3_container(parts): footer = pq.read_metadata(io.BytesIO(parts[0])).metadata assert footer[b"protspace_format_version"] == b"3" - assert pq.read_table(io.BytesIO(parts[5])).column_names == ["name", "data"] + assert _table(parts[5]).column_names == ["name", "data"] def test_manifest_declares_every_kind_the_reader_dispatches_on(parts): - manifest = json.loads(pq.read_metadata(io.BytesIO(parts[0])).metadata[MANIFEST_KEY]) + manifest = _manifest(parts[0]) assert manifest["idColumn"] == "protein_id" assert manifest["projections"] == [ @@ -84,8 +153,16 @@ def test_manifest_declares_every_kind_the_reader_dispatches_on(parts): assert manifest["columns"] == { "cath": {"kind": "multi", "sourceType": "string", "scores": True}, "go_bp": {"kind": "multi", "sourceType": "string", "evidence": True}, - "pfam": {"kind": "multi", "sourceType": "string", "scores": True}, + # The only column that declares both payload families, so the encoder and + # the reader cross that pair on real bytes exactly here. + "pfam": { + "kind": "multi", + "sourceType": "string", + "scores": True, + "evidence": True, + }, "kingdom": {"kind": "categorical", "sourceType": "string"}, + "reviewed": {"kind": "categorical", "sourceType": "string"}, "predicted_tm": {"kind": "categorical", "sourceType": "string"}, "length": {"kind": "numeric", "numericType": "int", "sourceType": "string"}, "hydrophobicity": { @@ -113,6 +190,7 @@ def test_part_one_is_the_physical_v3_schema(parts): "go_bp__count", "pfam__count", "kingdom", + "reviewed", "predicted_tm", "length", "hydrophobicity", @@ -124,6 +202,70 @@ def test_part_one_is_the_physical_v3_schema(parts): assert all(not field.nullable for field in schema) +def test_dictionaries_are_ordered_by_descending_frequency(parts, payloads): + """Dictionary order is the browser's legend order, so it is part of the contract. + + Both columns below are written so that descending frequency and first + occurrence DISAGREE. Nothing else in the fixture can tell the two orderings + apart, and a decoded comparison never could: the decoder re-joins the same + cells whichever order the codes are in. + """ + codes = _table(parts[0]).column("kingdom").to_pylist() + + # Cells: Archaea, Bacteria, Bacteria, , Bacteria, Eukaryota. + assert generator.ANNOTATION_CELLS["kingdom"][0] == "Archaea" + assert _labels(payloads, "kingdom") == ["Bacteria", "Archaea", "Eukaryota"] + assert codes == [1, 0, 0, -1, 0, 2] # -1 is the blank cell + + # Cells: False, True, True, False, True, True. + assert generator.ANNOTATION_CELLS["reviewed"][0] == "False" + assert _labels(payloads, "reviewed") == ["True", "False"] + assert _table(parts[0]).column("reviewed").to_pylist() == [1, 0, 0, 1, 0, 0] + + +def test_dictionary_label_lengths_are_utf8_bytes(payloads): + """The browser slices ``dict:`` by these lengths, so they must be bytes. + + Every other dictionary in the fixture is pure ASCII, where a byte length and + a character count are the same number and a reader that confused them would + still pass. This label is the one that separates them. + """ + labels = _labels(payloads, "pfam") + lengths = _i32(payloads, "dict:pfam:len") + + assert NON_ASCII_LABEL in labels + at = labels.index(NON_ASCII_LABEL) + assert lengths[at] == len(NON_ASCII_LABEL.encode()) > len(NON_ASCII_LABEL) + assert sum(lengths) == len(payloads["dict:pfam"]) + # The two labels after it would shift by the same two bytes if the lengths + # were counted in characters. + assert labels[at + 1 :] == ["none", "PF00003 (a|b)"] + + +def test_payloads_carry_scores_and_evidence_for_one_column(payloads): + """``pfam`` holds both families at once; every index is per hit, not per row.""" + # 8 hits: P2 has 2, P4 has 3 (one of them the folded ``none``), P5 has 3. + assert _i32(payloads, "csr:pfam") == [0, 1, 0, 2, 3, 4, 1, 0] + assert _i32(payloads, "score_count:pfam") == [2, 1, 1, 0, 0, 1, 1, 1] + assert np.frombuffer(payloads["scores:pfam"], "Uu;uV7{8~z?b^|@miAumEiJZXD;w)}?>bgj5ahI7yUwj_W2}D?nYX=nYw5qa zvJG&_ct8@#5=agjUlgB+COjY|z6i#sFD7afO;r9wcpyPvh=xQH6Y-qZbr=&*@}2wr zJKy)a=ljn0ntUgJ{@ufBilPXS#ZZjMc)W}OGg?M#$108n6c4;l@FE_n|MddIg{=4Q zb_3&`0gBUBvpJ@+*- zxuzxt^hnKuQ%InC0-Rof7UZ;CDuN38(vQ$#sdz?)RI^Qx`o6`LS zz%6vB^%lC{YC+((h4q<|D8B-jFAq-6r zK1g&!Q-gvKs5IqBrS{0!0d|arMv8K&SV$)Zj;0UA{POTfUpgsPW2IyBnUY+{_(X(E z2R5BZC)4xkq+CpgYVNjK^=26!#$hKg1&hIn+8bN;(z?>LPfASpr_rw#TkFON^w!q0 zX*pY(E=Y;Em`+U2%25%;J=}klWzUP~W7ee0M5hV? z7Yo@^U|EU=@;L0gJX}C!&(A#&rU<6xRAQizN{(TkTt}z^=O9eL7%_z)3>OIWDoXKP zs!$4@=EIiVBi-?)+`vovve77Y2G-L@<2gB++enYA zgZmZVM;7Ah{IeSIa8DnG=(N`cA_U2{3bqDuHwADf6&I7fVgW0#(*z8n&+XmhkvkY< za%c(FC(4MU#ygG}@cTq3y1L2dQvFl=w(E7Zbxp@SR(l!$#soS;qm`cA5wOp^22A_c zC}4H0fi`e}nE__XU_NIp3F>y$l6cJDJ%L8d%v?&TFG=-ji(1&@HBE^qA?jM zO7jo^^O|`DIL%h@5OBZ*@Mr@DGwmJ39v|NL{>PtIZtv4ODcU7kvF6_Mlm<{=8+DY< z$h99l@Cx&)=6Ch>KlX!{s3#5p1V@_&xGtkB0zcwVIz3K9JEs;Fm55ZVDu4%H z+BG=TW=t+y`_z~-rxf$Y;}Glx+kuBt_3SXuant#Q#i~*T+W-xYpg)CYY_1{~3&Ww% z!9A`UPM<#z8-hd0ad^S?1M-V|d;Pqf_j)KhW8gcq6wlFQh~g3MXFt%fZ_^%Ka~?m( rUeltBqOH@-T=slO{r)<|-c`dKS*x1Kb){lqwRb%bG*a`jy;N!f0 delta 1470 zcmZ7$OKclObY{K&B(a;g@osG5{2SNB0k79d6T38I{LPwt?GQVWz^p&b-ftdGGCe^X7-b z)yzxp4nY6_1-dxE37j~<=?S6XG*%)qw@?Ly>JhRrlqL&)@RtJ`z8nH(&00F2Ez2B_ zhnN(8hq(%GU_o`DO4#n>JE}H3&hpHsD*Oz47GB3mwhh0l`P({PZ`1M5>40UztXY*S8F|h8O9OttCtSO!K4fY= z4Sz%8zPsr*Zq$2gyZUCP@w-kZUGCfj{{75F@rhXpArvHl&=vtH7eL4-QQAjk($0v} zIg+dMKB>kw_rp_5v8{qEmNpi}=w>Qbd#zWSJX)3G5f=%yBNbP2S+^oDrpxpBOxROZWYx$Tv8gU3B14py zD!w5?6YuH5gdY_$<@st_EN!L=(s{`*Ri|QuTN6c8Ab(43^Qm$)pLTh~#YALBj(aw! z58_DLn5+)&NELxrEKE^49NJh6hjs`#q40>5*-t(cSf&-age&dmn7|}{y*9|>LKyvhk4EM zpW)bGze}lua?_L;$WQAw?jwDSmgM-*k{qYEaF)6~zS_96%Rwe^rmIsK-qY>1XS<%( zGpl3xYlnl`qgxy1229NJV>QRXt18o*TJ96o0Gaywz<;H}2hO}7css7ameFefnh4Ph z^v27kb}PbgP%L&}zzNlIvw|iUaUbU^ZF|J~o+z!}ty1Y_AAf zL#;sXEoCbt#=QpaJ?#VGj#zvs{0RH-XD)NZ?C!S9yQy>VH<&bO$?O|I!ygbDlW@AH z8S0?P$PU)l@`-@P}0l>mT6Or9}9H=4KH|SX1SdRmjMT-dc;7@RGu7PMz z5vZ(GvU8K&`E04MJ%wPejs-w%1I!}^?-YOtvM`AE#V)Izx|0GZ=%4oX^mU8pJyH+_ z6EpDVo>AQ3zHmjbSSrpg19rx1bw)VN(0{;Q_etg6L diff --git a/packages/core/src/components/data-loader/utils/v3-equivalence.test.ts b/packages/core/src/components/data-loader/utils/v3-equivalence.test.ts index f1cc660e..9a224458 100644 --- a/packages/core/src/components/data-loader/utils/v3-equivalence.test.ts +++ b/packages/core/src/components/data-loader/utils/v3-equivalence.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect } from 'vitest'; +import { describe, it, expect, vi, afterEach } from 'vitest'; import { readFileSync } from 'node:fs'; import { createParquetBundle, @@ -8,6 +8,7 @@ import { isCsrAnnotationData, isMultilabelAnnotationData, isNAValue, + materializeEatOverlay, NA_DEFAULT_COLOR, NA_VALUE, type VisualizationData, @@ -30,12 +31,17 @@ import { convertParquetToVisualizationDataOptimized } from './conversion'; * which cannot carry an interior zero-hit row, a numeric column or a 3D projection), so * the two are compared only on the values they genuinely share: `protein_ids[0..1]`, * P1/P2's `cath` and `go_bp`, and the `pca2` coordinates. + * + * The Python side asserts on the same bytes from the other end + * (`apps/protspace/tests/test_bundle_v3_fixture.py`), including the encoded part 1 and + * part 6 the decoder would otherwise hide. Change the fixture in the generator only, + * never by hand, and re-run both suites. */ const PROTEIN_IDS = ['P1', 'P2', 'P3', 'P4', 'P5', 'P6'] as const; /** Every categorical column of the fixture, in part-1 order. */ -const CATEGORICAL = ['cath', 'go_bp', 'pfam', 'kingdom', 'predicted_tm'] as const; +const CATEGORICAL = ['cath', 'go_bp', 'pfam', 'kingdom', 'reviewed', 'predicted_tm'] as const; /** * The label whose own text contains the `;` the v2 cell grammar reserves as the hit @@ -44,10 +50,32 @@ const CATEGORICAL = ['cath', 'go_bp', 'pfam', 'kingdom', 'predicted_tm'] as cons */ const CATH_ENCODED_SEMICOLON = 'G3DSA:1.10.10.10 (Ribosomal Protein L15; Chain: K; domain 2)'; +/** + * The fixture's one label outside ASCII. `readLabels` decodes a dictionary blob in one + * pass and slices it by character offset only while the blob is pure ASCII; this label + * is what forces the other branch, where each label is decoded from its own UTF-8 byte + * range. Python measures those lengths in bytes, so the two sides have to agree on them. + */ +const PFAM_NON_ASCII = 'PF00004 (β-lactamase, Nébuline)'; + +/** + * Which flat per-hit payload families each multi column actually carries. This is what + * an inserted-`__NA__` hit reports `null` for (see the round-trip suite below); a column + * that carries neither reports an empty array on every row. + */ +const PAYLOADS: Readonly> = { + cath: { scores: true, evidence: false }, + go_bp: { scores: false, evidence: true }, + pfam: { scores: true, evidence: true }, + kingdom: { scores: false, evidence: false }, + reviewed: { scores: false, evidence: false }, + predicted_tm: { scores: false, evidence: false }, +}; + /** * Rows whose only "hit" is the synthetic `__NA__` the reader inserts for an empty CSR - * row, per column that carries a score or evidence payload. These are exactly the rows - * where CSR and nested storage legitimately disagree (see the round-trip suite). + * row. These are exactly the rows where CSR and nested storage legitimately disagree + * (see the round-trip suite). */ const NA_ONLY_ROWS: Readonly> = { cath: ['P4'], @@ -88,6 +116,10 @@ function hitsOf(data: VisualizationData, key: string, proteinIndex: number): Hit const hitsByProtein = (data: VisualizationData, key: string): Record => Object.fromEntries(data.protein_ids.map((id, index) => [id, hitsOf(data, key, index)])); +afterEach(() => { + vi.restoreAllMocks(); +}); + describe('v3 golden fixture: the Python encoder and the browser reader agree', () => { it('reads the six-part container with both empty slots and no settings or statistics', async () => { const { data, settings } = await loadV3(); @@ -98,24 +130,88 @@ describe('v3 golden fixture: the Python encoder and the browser reader agree', ( expect(data.statistics).toBeUndefined(); }); + it('decodes every part-1 column straight into a typed array, with no plain-array fallback', async () => { + // The whole performance premise of v3: hyparquet hands back a typed array only for a + // REQUIRED, PLAIN, undictionaried column, and the reader logs (once) when it has to + // fall back to the ~4x slower element loop. Nothing else in the suite would notice: + // the fallback decodes correctly, so this is the only assertion that proves the + // Python writer really produced the physical shape the reader is optimised for. + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}); + + const { data } = await loadV3(); + + expect(data.protein_ids).toHaveLength(PROTEIN_IDS.length); + expect(warn).not.toHaveBeenCalled(); + }); + it('exposes exactly the declared annotations, with the EAT companion trio consumed', async () => { const { data } = await loadV3(); // `kingdom__pred_value/__pred_confidence/__pred_source` are declared in the manifest // and physically present in part 1; `normalizeEatCompanionColumns` must consume them - // so they never become three junk legend columns. + // so they never become three junk legend columns. In their place it synthesises one + // runtime-only numeric column for the confidence. expect(Object.keys(data.annotations)).toEqual([ 'cath', 'go_bp', 'pfam', 'kingdom', + 'reviewed', 'predicted_tm', 'length', 'hydrophobicity', + 'kingdom__eat_confidence', ]); - // Every prediction targets a protein whose curated `kingdom` is present, so the - // overlay yields no cells at all - and must not invent an empty record either. - expect(data.annotation_predicted).toBeUndefined(); + expect(data.annotations.kingdom__eat_confidence).toEqual({ + kind: 'numeric', + numericType: 'float', + values: [], + colors: [], + shapes: [], + runtime: { role: 'eat-confidence', baseAnnotation: 'kingdom' }, + }); + }); + + it('keeps only the prediction whose curated cell is missing, as a real predicted cell', async () => { + const { data } = await loadV3(); + + // The fixture carries three predictions. P2 and P5 have a curated `kingdom`, so the + // overlay must discard them; only P4's cell is blank, so only P4 gets a prediction. + expect(data.annotation_predicted).toEqual({ + kingdom: [ + null, + null, + null, + { value: 'Viruses', confidence: 0.5, source: 'P0A7B8' }, + null, + null, + ], + }); + // `Viruses` occurs nowhere in the curated column, so the overlay had to grow the + // legend by one prediction-only value, appended after the observed ones and before + // the synthetic NA, with a fresh palette colour. + expect(data.annotations.kingdom.values).toEqual([ + 'Bacteria', + 'Archaea', + 'Eukaryota', + 'Viruses', + NA_VALUE, + ]); + expect(data.annotations.kingdom.colors).toEqual([ + '#F3C300', + '#875692', + '#F38400', + '#A1CAF1', + NA_DEFAULT_COLOR, + ]); + // The source is not one of the six proteins, so no `sourceIndex` may be attached. + expect(data.annotation_predicted!.kingdom[3]).not.toHaveProperty('sourceIndex'); + + // Turning the overlay on moves P4 off the NA slot and onto `Viruses`, and leaves + // every curated row exactly where it was. + const overlaid = materializeEatOverlay(data, 'kingdom', true); + expect(Array.from(overlaid.annotation_data.kingdom as Int32Array)).toEqual([1, 0, 0, 3, 0, 2]); + expect(Array.from(data.annotation_data.kingdom as Int32Array)).toEqual([1, 0, 0, 4, 0, 2]); }); it('stores multi-valued columns as CSR and single-valued ones as flat codes', async () => { @@ -127,7 +223,7 @@ describe('v3 golden fixture: the Python encoder and the browser reader agree', ( expect(isMultilabelAnnotationData(storage), key).toBe(true); expect((storage as { length: number }).length).toBe(PROTEIN_IDS.length); } - for (const key of ['kingdom', 'predicted_tm'] as const) { + for (const key of ['kingdom', 'reviewed', 'predicted_tm'] as const) { expect(data.annotation_data[key], key).toBeInstanceOf(Int32Array); } }); @@ -202,36 +298,90 @@ describe('v3 golden fixture: the Python encoder and the browser reader agree', ( expect(data.annotations.pfam).toEqual({ kind: 'categorical', // The encoded '|' - the grammar's suffix separator - is decoded back into a label. - values: ['PF00001 (7tm;1)', 'PF00002', 'PF00003 (a|b)', NA_VALUE], - colors: ['#F3C300', '#875692', '#F38400', NA_DEFAULT_COLOR], - shapes: ['circle', 'circle', 'circle', 'circle'], + values: ['PF00001 (7tm;1)', 'PF00002', PFAM_NON_ASCII, 'PF00003 (a|b)', NA_VALUE], + colors: ['#F3C300', '#875692', '#F38400', '#A1CAF1', NA_DEFAULT_COLOR], + shapes: ['circle', 'circle', 'circle', 'circle', 'circle'], }); expect(hitsByProtein(data, 'pfam')).toEqual({ - P1: { labels: [NA_VALUE], scores: [null], evidence: [] }, + P1: { labels: [NA_VALUE], scores: [null], evidence: [null] }, // Two scores on one hit, one on the next: the score_count payload, not a 1:1 map. - P2: { labels: ['PF00001 (7tm;1)', 'PF00002'], scores: [[1e-10, 2.5], [0.5]], evidence: [] }, - P3: { labels: [NA_VALUE], scores: [null], evidence: [] }, + P2: { + labels: ['PF00001 (7tm;1)', 'PF00002'], + scores: [[1e-10, 2.5], [0.5]], + evidence: [null, null], + }, + P3: { labels: [NA_VALUE], scores: [null], evidence: [null] }, // P4 immediately follows the interior empty row: its score is what an off-by-one - // in the inserted-NA `hitEnd` would steal. - P4: { labels: ['PF00001 (7tm;1)'], scores: [[0.25]], evidence: [] }, - P5: { labels: ['PF00003 (a|b)'], scores: [[3]], evidence: [] }, - P6: { labels: [NA_VALUE], scores: [null], evidence: [] }, + // in the inserted-NA `hitEnd` would steal. It is also the one row that crosses a + // score and an evidence hit inside a single column, and it had a third hit + // spelled `none` that the reader folded away (see below). + P4: { + labels: ['PF00001 (7tm;1)', PFAM_NON_ASCII], + scores: [[0.25], null], + evidence: [null, 'IDA'], + }, + P5: { + labels: ['PF00003 (a|b)', 'PF00002', 'PF00001 (7tm;1)'], + // `62.0` and `2.3e-5` on the wire; both are just doubles by the time they land. + scores: [[3], [62], [2.3e-5]], + evidence: [null, null, null], + }, + P6: { labels: [NA_VALUE], scores: [null], evidence: [null] }, }); }); - it('decodes a plain categorical column with no NA slot at all', async () => { + it('folds a missing-value label out of a multi column, dropping its hit entirely', async () => { const { data } = await loadV3(); - expect(data.annotations.kingdom).toEqual({ + // Part 6 carries `none` as an ordinary fourth pfam label - the encoder is a faithful + // container - and the browser drops it from the dictionary AND drops P4's third hit + // with it, renumbering the codes, the score runs and the evidence codes in lockstep. + // Without this cell `dropFoldedHits` returns early on every column in the fixture. + expect(data.annotations.pfam.values).not.toContain('none'); + expect(hitsOf(data, 'pfam', PROTEIN_IDS.indexOf('P4')).labels).toHaveLength(2); + // Folding must not manufacture a second NA slot, and must not disturb the surviving + // labels' frequency order. + expect(data.annotations.pfam.values.filter(isNAValue)).toHaveLength(1); + }); + + it('reads a dictionary whose labels are not pure ASCII', async () => { + const { data } = await loadV3(); + + // Guard on the fixture itself: if this label ever loses its non-ASCII characters the + // reader silently goes back to slicing the whole blob by character offset, and the + // byte-range branch stops being covered by real encoder bytes. + expect(new TextEncoder().encode(PFAM_NON_ASCII).length).toBeGreaterThan(PFAM_NON_ASCII.length); + // Python measured this label's length in UTF-8 bytes; the browser has to slice the + // blob by the same measure or every later label in the dictionary shifts. + expect(data.annotations.pfam.values[2]).toBe(PFAM_NON_ASCII); + expect(data.annotations.pfam.values[3]).toBe('PF00003 (a|b)'); + }); + + it('orders a categorical dictionary by descending frequency, not first occurrence', async () => { + const { data } = await loadV3(); + + // `reviewed` is False, True, True, False, True, True: first occurrence would put + // `False` first, descending frequency puts `True` first. Dictionary order IS legend + // order and therefore colour assignment, so this is the assertion that fails if the + // encoder ever stops sorting. + expect(data.annotations.reviewed).toEqual({ kind: 'categorical', - values: ['Bacteria', 'Archaea', 'Eukaryota'], - colors: ['#F3C300', '#875692', '#F38400'], - shapes: ['circle', 'circle', 'circle'], + values: ['True', 'False'], + colors: ['#F3C300', '#875692'], + shapes: ['circle', 'circle'], }); - // Every row has a kingdom, so no synthetic category may be appended. - expect(data.annotations.kingdom.values.some(isNAValue)).toBe(false); - expect(Array.from(data.annotation_data.kingdom as Int32Array)).toEqual([0, 1, 0, 2, 0, 1]); + expect(Array.from(data.annotation_data.reviewed as Int32Array)).toEqual([1, 0, 0, 1, 0, 0]); + // Every row has a value, so no synthetic category may be appended. + expect(data.annotations.reviewed.values.some(isNAValue)).toBe(false); + + // Same divergence on `kingdom`, whose curated values are Archaea, Bacteria, + // Bacteria, , Bacteria, Eukaryota. + expect(data.annotations.kingdom.values.slice(0, 3)).toEqual([ + 'Bacteria', + 'Archaea', + 'Eukaryota', + ]); }); it('folds every missing-value spelling in one dictionary into a single NA slot', async () => { @@ -265,16 +415,22 @@ describe('v3 golden fixture: the Python encoder and the browser reader agree', ( colors: [], shapes: [], }); - expect(data.annotations.hydrophobicity).toMatchObject({ + expect(data.annotations.hydrophobicity).toEqual({ kind: 'numeric', numericType: 'float', + values: [], + colors: [], + shapes: [], }); expect(data.numeric_annotation_data).toEqual({ length: [120, null, 340, 0, -15, 1024], hydrophobicity: [0.5, -1.25, null, 3, 0.001, 42], + // Not a wire column: synthesised from the EAT confidence companion. + kingdom__eat_confidence: [null, null, null, 0.5, null, null], }); - // A numeric column carries no categorical storage to bin by code. - expect(data.annotation_data.length).toBeUndefined(); + // A numeric column carries no categorical storage to bin by code. Spelled as the + // whole key set so `annotation_data['length']` cannot be mistaken for an array length. + expect(Object.keys(data.annotation_data)).toEqual([...CATEGORICAL]); }); it('interleaves the wide axis columns into a 2D and a 3D projection', async () => { @@ -285,11 +441,18 @@ describe('v3 golden fixture: the Python encoder and the browser reader agree', ( expect(pca2.dimension).toBe(2); expect(Array.from(pca2.data)).toEqual([0, 0, 1, 1, 2.5, -3.5, -4, 0.25, 5, 5, -1.5, 2]); - expect(pca2.metadata).toMatchObject({ components: 2, dimension: 2, dimensions: 2 }); + expect(pca2.metadata).toEqual({ components: 2, dimension: 2, dimensions: 2, source: '' }); expect(umap3.dimension).toBe(3); - expect(Array.from(umap3.data)).toEqual(Array.from({ length: 18 }, (_, index) => index / 4)); - expect(umap3.metadata).toMatchObject({ n_neighbors: 15, dimension: 3, dimensions: 3 }); + expect(Array.from(umap3.data)).toEqual([ + ...Array.from({ length: 15 }, (_, index) => index / 4), + // P6 has no umap3 row at all. The encoder writes 0.0 for it and the browser leaves + // its zero-initialised slot untouched, so both put it at the origin. + 0, + 0, + 0, + ]); + expect(umap3.metadata).toEqual({ n_neighbors: 15, dimension: 3, dimensions: 3, source: '' }); }); }); @@ -314,7 +477,8 @@ describe('v3 -> v2 export round trip: CSR storage is interchangeable with nested cath: `${cathSemicolon}|50.2;G3DSA:6.20.10.10|60.5`, go_bp: 'apoptotic process|IDA', pfam: null, - kingdom: 'Bacteria', + kingdom: 'Archaea', + reviewed: 'False', // Documented non-identity: `none` is a MISSING_VALUE_TOKEN, so it was folded to // `__NA__` on read and goes back out as NULL, not as the literal word. predicted_tm: null, @@ -323,7 +487,8 @@ describe('v3 -> v2 export round trip: CSR storage is interchangeable with nested cath: '6.20.10.10', go_bp: null, pfam: 'PF00001 (7tm%3B1)|1e-10,2.5;PF00002|0.5', - kingdom: 'Archaea', + kingdom: 'Bacteria', + reviewed: 'True', predicted_tm: null, }, P3: { @@ -331,20 +496,32 @@ describe('v3 -> v2 export round trip: CSR storage is interchangeable with nested go_bp: 'apoptotic process|IDA;protein folding|ECO:0000269', pfam: null, kingdom: 'Bacteria', + reviewed: 'True', predicted_tm: 'TM helix', }, P4: { cath: null, go_bp: 'protein folding|IEA', - pfam: 'PF00001 (7tm%3B1)|0.25', - kingdom: 'Eukaryota', + // A score and an evidence code side by side in one cell, and the third hit - + // the one spelled `none` - gone, the same way the browser drops a folded label + // out of a v2 cell. + pfam: `PF00001 (7tm%3B1)|0.25;${PFAM_NON_ASCII}|IDA`, + // P4's curated cell was blank and now carries a prediction, so the base column + // goes back out NULL and the label rides in the companion trio instead. + kingdom: null, + reviewed: 'False', predicted_tm: null, }, P5: { cath: `${cathSemicolon}|1e-200`, go_bp: null, - pfam: 'PF00003 (a%7Cb)|3', + // Both documented score re-spellings. `62.0` loses its trailing `.0` on both + // sides; `2.3e-5` is where the two languages genuinely differ - Python's + // `read_tables` writes `2.3e-05`, `String(2.3e-5)` here writes `0.000023`. The + // double is identical, only the spelling is not. + pfam: 'PF00003 (a%7Cb)|3;PF00002|62;PF00001 (7tm%3B1)|0.000023', kingdom: 'Bacteria', + reviewed: 'True', // The other missing-value spelling in the same column, same treatment. predicted_tm: null, }, @@ -352,10 +529,20 @@ describe('v3 -> v2 export round trip: CSR storage is interchangeable with nested cath: '6.20.10.10', go_bp: 'apoptotic process|EXP', pfam: null, - kingdom: 'Archaea', + kingdom: 'Eukaryota', + reviewed: 'True', predicted_tm: 'TM helix', }, }); + // The prediction survives the export as the companion trio it arrived in. + expect([...extraction.annotationsById.values()].map((row) => row.kingdom__pred_value)).toEqual([ + null, + null, + null, + 'Viruses', + null, + null, + ]); // The in-memory sentinel is never written as a literal 6-char category. for (const row of extraction.annotationsById.values()) { expect(Object.values(row)).not.toContain(NA_VALUE); @@ -371,7 +558,7 @@ describe('v3 -> v2 export round trip: CSR storage is interchangeable with nested // they have to survive a shape change that reorders nothing. expect(reloaded.annotations).toEqual(v3.annotations); expect(reloaded.numeric_annotation_data).toEqual(v3.numeric_annotation_data); - expect(reloaded.annotation_predicted).toBeUndefined(); + expect(reloaded.annotation_predicted).toEqual(v3.annotation_predicted); expect( reloaded.projections.map(({ name, dimension, data }) => ({ name, @@ -406,16 +593,20 @@ describe('v3 -> v2 export round trip: CSR storage is interchangeable with nested expect(from2.labels, where).toEqual(from3.labels); - // The ONE documented non-identity. An empty CSR row owns no hit slot, so the - // reader inserts a synthetic `__NA__` hit for it - and the flat score/evidence - // payloads are numbered by hit, so that inserted hit reports itself as `null`. - // Nested storage has no hit there at all and reports nothing. Asserted as the - // exact rows it applies to rather than by relaxing the comparison. + // The documented non-identity, and it applies to BOTH payload families. An empty + // CSR row owns no hit slot, so the reader inserts a synthetic `__NA__` hit for it + // - and the flat score and evidence payloads are numbered by hit, so that + // inserted hit reports itself as `null` in whichever families the column carries. + // Nested storage has no hit there at all and reports nothing. Left as it is on + // purpose: none of the four consumers (tooltip, export, legend, statistics + // popover) distinguishes `[null]` from `[]`, and the flat shape is what keeps the + // score and evidence indices aligned with `getProteinAnnotationIndices`. Asserted + // as the exact rows it applies to rather than by relaxing the comparison. if (NA_ONLY_ROWS[key]?.includes(id)) { + const { scores, evidence } = PAYLOADS[key]; expect(from3.labels, where).toEqual([NA_VALUE]); - const scored = key !== 'go_bp'; - expect(from3.scores, where).toEqual(scored ? [null] : []); - expect(from3.evidence, where).toEqual(scored ? [] : [null]); + expect(from3.scores, where).toEqual(scores ? [null] : []); + expect(from3.evidence, where).toEqual(evidence ? [null] : []); expect(from2.scores, where).toEqual([]); expect(from2.evidence, where).toEqual([]); continue; @@ -433,7 +624,7 @@ describe('v2 and v3 fixtures agree on the values they share', () => { const { data: v3 } = await loadV3(); const v2 = await loadLegacy(fixture('v2-sample.parquetbundle')); - // The v3 fixture is a superset: 6 proteins to the v2 sample's 2, and 7 columns to + // The v3 fixture is a superset: 6 proteins to the v2 sample's 2, and 8 columns to // its 2. Only the shared prefix is comparable. expect(v2.protein_ids).toEqual(v3.protein_ids.slice(0, 2));