diff --git a/apps/protspace/CLAUDE.md b/apps/protspace/CLAUDE.md index d9c3fbec4..c359c3c5e 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/apps/protspace/scripts/generate_toxprot_demo.py b/apps/protspace/scripts/generate_toxprot_demo.py index 0a6cdfbc6..50dcaeeea 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/generate_v3_fixture.py b/apps/protspace/scripts/generate_v3_fixture.py new file mode 100644 index 000000000..0352987f2 --- /dev/null +++ b/apps/protspace/scripts/generate_v3_fixture.py @@ -0,0 +1,235 @@ +#!/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 seven more columns cover the v3 +paths a two-row two-column table cannot reach: + +* 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`` 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, 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. 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. + +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 +import pyarrow.compute as pc + +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})" + +# 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": [ + 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. + # 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", + "", + 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", + "", + ], + # 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"], + "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", "Viruses", 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) + + 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), + coordinates, + ] + + +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/scripts/inspect_bundle.py b/apps/protspace/scripts/inspect_bundle.py index c601272de..2791b844c 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 2a1b6bd08..81e3ba4e9 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 1ebd7eb68..74e67ee6a 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/annotations/encoding.py b/apps/protspace/src/protspace/data/annotations/encoding.py index 889acedcd..c5760696f 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/src/protspace/data/io/__init__.py b/apps/protspace/src/protspace/data/io/__init__.py index 0c0b0feec..9b30ac18a 100644 --- a/apps/protspace/src/protspace/data/io/__init__.py +++ b/apps/protspace/src/protspace/data/io/__init__.py @@ -12,9 +12,11 @@ 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, + read_tables, replace_settings_in_bundle, write_bundle, ) @@ -31,7 +33,9 @@ "PARQUET_BUNDLE_DELIMITER", "extract_bundle_to_dir", "read_bundle", + "read_settings_from_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 2c17504ed..57601c3de 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,59 @@ 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). - 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``). + 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``. """ - with open(bundle_path, "rb") as f: - parts = f.read().split(PARQUET_BUNDLE_DELIMITER) + 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): + return 1 - if len(parts) < 3 or len(parts) > 5: - raise ValueError(f"Expected 3 to 5 parts in parquetbundle, found {len(parts)}") + +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. + """ + 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)}") + + 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 +147,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 +226,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 +246,38 @@ 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_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.""" - _, _, statistics = _parse_bundle(bundle_path) - return statistics + return _parse_bundle(bundle_path)[2] def write_bundle( @@ -162,40 +288,50 @@ 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. + + **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). 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 +342,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 +359,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/src/protspace/data/io/bundle_v3.py b/apps/protspace/src/protspace/data/io/bundle_v3.py new file mode 100644 index 000000000..b3412ed79 --- /dev/null +++ b/apps/protspace/src/protspace/data/io/bundle_v3.py @@ -0,0 +1,900 @@ +"""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 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 +``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 +import logging +from typing import Any + +import numpy as np +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, + encode_field, + migrate_legacy_annotation_table, + read_format_version, + stamp_format_version, +) + +logger = logging.getLogger(__name__) + +CONTAINER_VERSION = 3 +MANIFEST_KEY = b"protspace_v3_manifest" + +#: 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 +#: ``|`` 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. 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, 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 +#: 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" + +#: ``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 +# --------------------------------------------------------------------------- # + + +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 _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 + 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 _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.""" + token = pc.is_in(pc.utf8_lower(trimmed), value_set=pa.array(sorted(MISSING_TOKENS))) + return _blank_mask(trimmed) | np.asarray(pc.fill_null(token, False)) + + +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) -> 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) + + +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::len`` int32 per-label byte lengths.""" + encoded = [label.encode("utf-8") for label in labels] + lengths = np.array([len(b) for b in encoded], dtype=np.int64) + return [ + (f"dict:{name}", b"".join(encoded)), + (f"dict:{name}:len", _counts_i32(lengths, f"dict:{name}").tobytes()), + ] + + +def _split_last_pipe(hits: pa.Array) -> 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 ``__count`` per-row CSR hit counts; the + caller picks the physical column name from ``manifest_entry["kind"]``. + """ + 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 + # 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)] + 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, + "sourceType": source_type, + } + return entry, pa.array(values, type=pa.float64()), [] + + strings = _as_string(arr) + trimmed = pc.utf8_trim_whitespace(strings) + 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(): + values = _parse_floats(trimmed, ~missing) + 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 --------------------------------- # + # ``_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 = ~_blank_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) + # ``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) + 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 + + row_counts = _counts_i32(per_row, f"column '{name}' hits") + payloads.append((f"csr:{name}", codes.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}" + ) + + name_column = projections_data.column("projection_name") + # The browser derives the projection set, order and dimension from the data + # rows alone (``conversion.ts:1163-1196``), so a metadata-only projection + # would be an all-zero one there and a data-only projection would silently + # vanish here. All five shipped datasets agree; refuse the ones that do not. + in_data = set(pc.unique(name_column).to_pylist()) + if in_data != set(names): + raise ValueError( + "projections_metadata and projections_data disagree on the projection " + f"set: metadata-only {sorted(set(names) - in_data, key=str)}, " + f"data-only {sorted(in_data - set(names), key=str)}" + ) + + 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 + + num_rows = len(protein_ids) + columns: dict[str, pa.Array] = {} + manifest: list[dict[str, Any]] = [] + + for name, declared in zip(names, dimensions, strict=True): + rows = projections_data.filter(pc.equal(name_column, pa.scalar(name))) + identifiers = rows.column("identifier") + found = pc.index_in(identifiers, value_set=protein_ids) + if found.null_count: + unknown = identifiers.filter(pc.is_null(found)).to_pylist() + raise ValueError( + f"projection '{name}' references identifier(s) absent from the " + f"annotations table: {sorted(set(unknown))[:5]}" + ) + positions = np.asarray(found.combine_chunks()) + if np.unique(positions).size != positions.size: + repeated = np.flatnonzero(np.bincount(positions, minlength=num_rows) > 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 = 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) + ) + 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]: + # 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) + ) + 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. + + ``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( + (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}__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 " + "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_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()), + "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), + ) + + +# --------------------------------------------------------------------------- # +# 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). + + 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) + 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. + + 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. + + 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) + + +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, 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( + 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, ",", f"payload 'score_count:{name}'"), + 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, ";", f"column '{name}__count'") + + +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: + + * 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``; + * 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``; + * 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. + """ + 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/src/protspace/utils/add_annotation_style.py b/apps/protspace/src/protspace/utils/add_annotation_style.py index 0bbc4df31..9cbc98cac 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/src/protspace/utils/arrow_reader.py b/apps/protspace/src/protspace/utils/arrow_reader.py index 32f88e67b..b8d8f537d 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_overlay.py b/apps/protspace/tests/test_bundle_overlay.py index 0b90bc6ff..54a456e91 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_settings.py b/apps/protspace/tests/test_bundle_settings.py index dc305a50b..66ea40721 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 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 000000000..bf32967c3 --- /dev/null +++ b/apps/protspace/tests/test_bundle_v3_container.py @@ -0,0 +1,424 @@ +"""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 +import logging +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_settings_from_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_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" + 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_v3_decode.py b/apps/protspace/tests/test_bundle_v3_decode.py new file mode 100644 index 000000000..4a7f014da --- /dev/null +++ b/apps/protspace/tests/test_bundle_v3_decode.py @@ -0,0 +1,567 @@ +"""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, _flat, 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_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", "NA", "n/a", "None", "__NA__", ""] + + +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_float64(): + table = annotations_table(col=["A|0.5700", "A|1", "A|0.1", "A|1e-10,2.5"]) + # 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"] + + +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( + 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) + + +# --------------------------------------------------------------------------- # +# 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 +# --------------------------------------------------------------------------- # + + +@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() 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 000000000..608a1701a --- /dev/null +++ b/apps/protspace/tests/test_bundle_v3_encode.py @@ -0,0 +1,574 @@ +"""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, + migrate_legacy_annotation_table, + read_format_version, + 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]: + """Rebuild the labels the way the reader does: prefix-sum the byte lengths.""" + blob = payloads[f"dict:{column}"] + ends = np.cumsum(np.frombuffer(payloads[f"dict:{column}:len"], ")`` 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 +# --------------------------------------------------------------------------- # + + +def test_code_order_is_frequency_then_first_occurrence(): + # counts: rare=1, common=3, tie_b=2, tie_a=2 (tie_b appears first). + table = make_annotations( + col=["rare", "common", "tie_b", "tie_a", "common", "tie_b", "tie_a", "common"] + ) + parts = encode(table) + assert labels_of(payloads_of(parts[3]), "col") == [ + "common", + "tie_b", + "tie_a", + "rare", + ] + assert read(parts[0]).column("col").to_pylist() == [3, 0, 1, 2, 0, 1, 2, 0] + + +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", + "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(): + table = make_annotations(pfam=["PF1|1e-10,2.5;PF2|0.5", "PF1|1.0", "", "PF3"]) + parts = encode(table) + payloads = payloads_of(parts[3]) + assert labels_of(payloads, "pfam") == ["PF1", "PF2", "PF3"] + assert read(parts[0]).column("pfam__count").to_pylist() == [2, 1, 0, 1] + assert list(np.frombuffer(payloads["csr:pfam"], " 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 + # 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 _table(parts[5]).column_names == ["name", "data"] + + +def test_manifest_declares_every_kind_the_reader_dispatches_on(parts): + manifest = _manifest(parts[0]) + + 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}, + # 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": { + "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", + "reviewed", + "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_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"], " 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 c2b6156ec..1af769724 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,11 @@ 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). 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"] == "" diff --git a/apps/web/src/explore/dataset-controller.eat.test.ts b/apps/web/src/explore/dataset-controller.eat.test.ts index 55efe49a2..85d7bd11a 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 new file mode 100644 index 000000000..4a774a761 --- /dev/null +++ b/apps/web/src/explore/dataset-controller.persistence.test.ts @@ -0,0 +1,181 @@ +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(), + saveLastImportedFileMetadata: vi.fn(), + saveLastImportedFileData: 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, + saveLastImportedFileMetadata: mocks.saveLastImportedFileMetadata, + saveLastImportedFileData: mocks.saveLastImportedFileData, +})); + +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.saveLastImportedFileMetadata.mockResolvedValue(undefined); + mocks.saveLastImportedFileData.mockResolvedValue(undefined); + }); + + /** 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(); + + // 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(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.markLastLoadStatus).toHaveBeenCalledWith('success'); + expect(mocks.resolvePendingLoadFinalization).toHaveBeenCalledWith(3); + }); + + 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.saveLastImportedFileData.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(); + }); + + 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 1c846c716..099f7701a 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'; @@ -116,21 +120,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); @@ -138,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') { @@ -236,6 +245,15 @@ export function createDatasetController({ viewController.applyLatestViewForDatasetLoad(data); + // 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 { if (loadMeta.kind === 'user' || loadMeta.kind === 'opfs') { await markLastLoadStatus('success'); diff --git a/apps/web/src/explore/opfs-dataset-store.test.ts b/apps/web/src/explore/opfs-dataset-store.test.ts index 356043762..392085479 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 746e961c1..84669c325 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.'); diff --git a/docs/guide/data-format.md b/docs/guide/data-format.md index b5b86fc69..1ba7b9df1 100644 --- a/docs/guide/data-format.md +++ b/docs/guide/data-format.md @@ -4,7 +4,27 @@ 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 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) ``` .parquetbundle file @@ -23,20 +43,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 +181,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 +239,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 +274,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 +303,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 +310,229 @@ 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. +- 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. +- 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. 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 +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 16 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. + +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 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 + +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. + +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 + +`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 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 +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. diff --git a/packages/core/src/components/data-loader/data-loader.ts b/packages/core/src/components/data-loader/data-loader.ts index d0f16d94b..ec72ada4f 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/decode-transferables.test.ts b/packages/core/src/components/data-loader/decode-transferables.test.ts new file mode 100644 index 000000000..c4ba2c9c2 --- /dev/null +++ b/packages/core/src/components/data-loader/decode-transferables.test.ts @@ -0,0 +1,108 @@ +import { describe, it, expect } from 'vitest'; +import { isCsrAnnotationData, 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'] }, + }, + }, + }; +} + +/** 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(); + 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: the clone holds the bytes and every source is detached', () => { + const { data } = csrDataset(); + const transfer = collectTransferables(data); + 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', () => { + 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 000000000..cf1be8465 --- /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 b5504a476..db45f6e39 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), 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 000000000..5cb32e37d Binary files /dev/null and b/packages/core/src/components/data-loader/utils/__fixtures__/v3-sample.parquetbundle differ 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 479b85e34..7658497a5 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/bundle-v3.test.ts b/packages/core/src/components/data-loader/utils/bundle-v3.test.ts new file mode 100644 index 000000000..f6e5d300b --- /dev/null +++ b/packages/core/src/components/data-loader/utils/bundle-v3.test.ts @@ -0,0 +1,718 @@ +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, + getProteinAnnotationIndices, + getProteinEvidence, + getProteinScores, + isCsrAnnotationData, + NA_DEFAULT_COLOR, + NA_VALUE, + type CsrAnnotationData, + type VisualizationData, +} from '@protspace/utils'; +import { decodeParquetBundle, extractRowsFromParquetBundle } from './bundle'; +import { readV3Bundle } from './bundle-v3'; +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); +const f64 = (...values: number[]) => new Uint8Array(new Float64Array(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': 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), + '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, + ), + ); + +/** + * 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 | Float64Array)[] => [ + ...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], + ); + +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('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()); + + 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/, + ], + [ + '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', + { + ...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': 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/, + ); + }); + + 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 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( + /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/, + ); + }); + }); + + // `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. + 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('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( + 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 x 2 CSR (end + codes) + scores (hitEnd + + // values) + evidence codes. + expect(transfer).toHaveLength(10); + + 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 new file mode 100644 index 000000000..1018cbc53 --- /dev/null +++ b/packages/core/src/components/data-loader/utils/bundle-v3.ts @@ -0,0 +1,812 @@ +/** + * 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. + * + * 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 + * 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, + normalizeMissingValue, + type Annotation, + type AnnotationData, + type BundleSettings, + type CsrEvidence, + type CsrScores, + type Projection, + type VisualizationData, +} from '@protspace/utils'; +import { assertValidParquetMagic, DEFAULT_VALIDATION_LIMITS } 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; + +// 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'; + +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) 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) types.set(field.name, field.type); + } + 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); +} + +/** + * 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 = leafColumnTypes(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)}"`, + ); + } + if (name === idColumn) { + throw new Error(`v3 manifest declares idColumn "${name}" as an annotation column too`); + } + const physical = physicalColumn(name, kind); + 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 } : {}), + ...(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); + } +} + +/** + * 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, + 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), + ); + } + + const onPlainArray = plainArrayReporter(); + + 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 = leafColumnTypes(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) { + const onPlainArray = plainArrayReporter(); + await parquetRead({ + file: part, + metadata, + columns: [...axisTargets.keys()], + 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; + } + }, + }); + } + + 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); + // 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; +} + +/** + * 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; readonly BYTES_PER_ELEMENT: number }, +): T { + const bytes = payloads.get(name); + if (!bytes) throw new Error(`v3 bundle is missing the "${name}" payload`); + 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); +} + +/** + * 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; +} + +/** + * 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); + 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; +} + +/** + * 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: V3Scores | 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: V3Scores | 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}`, Float64Array); + 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`, + ); + } + 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 }; +} + +/** + * 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. + * + * 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); + // 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[]; + + 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); + // 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] >= encodedLabelCount || codes[i] < -1) { + throw new Error( + `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( + dropFoldedHits( + readCsrColumn( + name, + column, + stored as Int32Array, + encodedLabelCount, + payloads, + readEvidenceDict, + ), + remap, + ), + 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 faee81c4e..d9ad5b0e0 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 88a85dd53..be7dd68d3 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; +} - // 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); +/** + * 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 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,19 @@ 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; + // 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) + : {}; // 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 +303,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 +340,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 +376,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.test.ts b/packages/core/src/components/data-loader/utils/conversion.test.ts index 13dc3a0da..72ba05cd0 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,135 @@ 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']`. 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 → [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, 2, 4), + codes: Int32Array.of(0, 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 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 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]); // 'B', renumbered + expect(getProteinAnnotationIndices(rows, 3)).toEqual([0, 1]); + expect(Array.from(rows.end)).toEqual([0, 1, 2, 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, 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', () => { + 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 0ddcf0a3f..4ebc351f8 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, @@ -131,7 +134,7 @@ function* valuesForColumn(rows: Rows, column: string): Iterable { } } -function createNumericAnnotation( +export function createNumericAnnotation( numericType: 'int' | 'float', runtime?: Annotation['runtime'], ): Annotation { @@ -222,6 +225,34 @@ function remapCategoricalStorage( } return { kind: 'sparse-multi', base, overrides, length: base.length }; } + if (isCsrAnnotationData(source)) { + // 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; + } + 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)); } @@ -230,7 +261,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 +279,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 +291,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 +318,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 +412,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, }; @@ -407,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` @@ -592,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>(); @@ -692,6 +757,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; } @@ -732,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; @@ -1657,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); 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 000000000..9a2244588 --- /dev/null +++ b/packages/core/src/components/data-loader/utils/v3-equivalence.test.ts @@ -0,0 +1,650 @@ +import { describe, it, expect, vi, afterEach } from 'vitest'; +import { readFileSync } from 'node:fs'; +import { + createParquetBundle, + getProteinAnnotationIndices, + getProteinEvidence, + getProteinScores, + isCsrAnnotationData, + isMultilabelAnnotationData, + isNAValue, + materializeEatOverlay, + 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. + * + * 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', 'reviewed', '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)'; + +/** + * 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. 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)])); + +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(); + + 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('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. 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', + ]); + 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 () => { + 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', 'reviewed', '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', 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: [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: [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. 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('folds a missing-value label out of a multi column, dropping its hit entirely', async () => { + const { data } = await loadV3(); + + // 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: ['True', 'False'], + colors: ['#F3C300', '#875692'], + shapes: ['circle', 'circle'], + }); + 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 () => { + 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).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. 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 () => { + 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).toEqual({ components: 2, dimension: 2, dimensions: 2, source: '' }); + + expect(umap3.dimension).toBe(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: '' }); + }); +}); + +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: '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, + }, + P2: { + cath: '6.20.10.10', + go_bp: null, + pfam: 'PF00001 (7tm%3B1)|1e-10,2.5;PF00002|0.5', + kingdom: 'Bacteria', + reviewed: 'True', + predicted_tm: null, + }, + P3: { + cath: 'G3DSA:6.20.10.10|123456789', + 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', + // 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, + // 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, + }, + P6: { + cath: '6.20.10.10', + go_bp: 'apoptotic process|EXP', + pfam: null, + 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); + } + }); + + 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).toEqual(v3.annotation_predicted); + 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 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]); + 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; + } + + 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 8 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))); + }); +}); 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 34840786f..000000000 --- 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 8da7dae86..000000000 --- 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 5223e96ac..be1b4e336 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 d33d382c6..b05aa9618 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 e09b5c64c..e518bc385 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,10 +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(); - if ( - !this.annotationData?.values?.length || - (!isNumericAnnotation && !this.annotationValues?.length) - ) { + // 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 (!isNumericAnnotation && !(frequencies?.size ?? this.annotationValues?.length)) { this._legendItems = []; return; } @@ -1801,9 +1856,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 +1877,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/interaction/quadtree-index.parity.test.ts b/packages/core/src/components/scatter-plot/interaction/quadtree-index.parity.test.ts new file mode 100644 index 000000000..47bf937b4 --- /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 4b759f5d8..a7c87e607 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,57 @@ -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 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. */ +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 +63,7 @@ export class QuadtreeIndex { rebuild(pd: PlotData, slots: ArrayLike) { if (!this.scales || slots.length === 0) { - this.qt = null; + this.clear(); return; } @@ -34,61 +72,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 +238,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, 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 9716d3232..f694ff298 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)', () => { 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 6213ec80a..a1e92ff56 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 d729ea0a3..27e32b0be 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/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 ef196235f..7f86f6663 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,115 @@ 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('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; + 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 89706b359..07542bc76 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); } diff --git a/packages/utils/src/parquet/bundle-writer.ts b/packages/utils/src/parquet/bundle-writer.ts index 48f2243d4..f0e57ca02 100644 --- a/packages/utils/src/parquet/bundle-writer.ts +++ b/packages/utils/src/parquet/bundle-writer.ts @@ -28,10 +28,15 @@ 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'; +/** 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'; @@ -133,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++) { @@ -143,6 +158,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 = 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]; @@ -152,8 +174,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/storage/data-hash.test.ts b/packages/utils/src/storage/data-hash.test.ts index 84e2eb394..b365a93ac 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 5f768c396..f3badf6cc 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,62 @@ 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. + * 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']; + 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; } diff --git a/packages/utils/src/types.ts b/packages/utils/src/types.ts index 6a38ea7bc..4cf76ae8b 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,48 @@ 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; + /** + * 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; +} + +/** + * 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 +193,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 0a6c5057e..3d183fc0c 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 9d506d8d0..39c67852b 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/eat-overlay.test.ts b/packages/utils/src/visualization/eat-overlay.test.ts index 71f475031..68928abda 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 { @@ -14,6 +15,7 @@ import { normalizeReliability, parseEatCompanionColumn, } from './eat-overlay'; +import { getProteinEvidence, getProteinScores } from './plot-data-accessors'; function createData(): VisualizationData { return { @@ -169,3 +171,157 @@ 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); + }); + + // 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: Float64Array.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 * Float64Array.BYTES_PER_ELEMENT, + ); + }); + + 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 = [ + { 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 10bad58db..fe8c73108 100644 --- a/packages/utils/src/visualization/eat-overlay.ts +++ b/packages/utils/src/visualization/eat-overlay.ts @@ -1,7 +1,15 @@ -import type { AnnotationData, PredictedCell, VisualizationData } from '../types.js'; +import type { + AnnotationData, + CsrAnnotationData, + CsrEvidence, + CsrScores, + PredictedCell, + VisualizationData, +} from '../types.js'; import { getFirstAnnotationIndex, getProteinAnnotationIndices, + isCsrAnnotationData, isSparseMultiValueAnnotationData, } from './annotation-data-access.js'; import { isNAValue } from './missing-values.js'; @@ -195,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 Float64Array(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 { @@ -277,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: { diff --git a/packages/utils/src/visualization/plot-data-accessors.test.ts b/packages/utils/src/visualization/plot-data-accessors.test.ts index ffbfef879..0a2f0fbe9 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: Float64Array.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 4dd49070c..00fbe3b6a 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; } /** diff --git a/packages/utils/src/visualization/slice-visualization-data.test.ts b/packages/utils/src/visualization/slice-visualization-data.test.ts index 0a3b604c1..46b207d34 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: Float64Array.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 c256a0625..bbf4f51d8 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 Float64Array(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), }; }