diff --git a/CHANGELOG.md b/CHANGELOG.md index c2e7f2232..c03a4a9cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -54,6 +54,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 passed directly to native community detection, and the recovered components are unchanged. ### Fixed +- `Edgelist.to_record_batches` renames legacy `marker1` and `marker2` columns to `marker_1` and `marker_2`, matching `to_polars`, `to_df`, and `iterator`. - `coarsened_pmds_layout` sizes PMDS pivots from the full graph when Leiden yields too few communities, so a valid low `pivots` no longer fails `pmds_layout`'s `0.2 * n` lower bound. diff --git a/src/pixelator/pna/pixeldataset/edgelist.py b/src/pixelator/pna/pixeldataset/edgelist.py index 3bb7e0dab..11a829d7f 100644 --- a/src/pixelator/pna/pixeldataset/edgelist.py +++ b/src/pixelator/pna/pixeldataset/edgelist.py @@ -11,8 +11,9 @@ import polars as pl import pyarrow as pa -from pixelator.pna.pixeldataset.io import PixelDataViewer, QueryBuilder +from pixelator.pna.pixeldataset.io import PixelDataViewer, Query, QueryBuilder from pixelator.pna.pixeldataset.io.anndata_helper import AnnDataHelper +from pixelator.pna.pixeldataset.legacy import LegacyMarkerQueryAdapter from pixelator.pna.pixeldataset.types import Component from pixelator.pna.utils import normalize_input_to_list, normalize_input_to_set @@ -56,9 +57,11 @@ def components(self) -> set[str]: ) return set(adata.obs.index.to_list()) - def _handle_backwards_compatibility(self, df: pl.LazyFrame) -> pl.LazyFrame: - # Handle legacy marker names - return df.rename({"marker1": "marker_1", "marker2": "marker_2"}, strict=False) + def _edgelist_query( + self, adapter: LegacyMarkerQueryAdapter, components: list[str] | None + ) -> Query: + """Build the edgelist query, adapted for a legacy marker schema when needed.""" + return adapter.adapt(self._query_builder.edgelist_query(components)) def __len__(self) -> int: """Get the number of edges in the edgelist.""" @@ -74,12 +77,11 @@ def is_empty(self) -> bool: def to_df(self) -> pd.DataFrame: """Get the edgelist as a pandas DataFrame.""" - query = self._query_builder.edgelist_query( - normalize_input_to_list(self.components) - ) + components = normalize_input_to_list(self.components) with self._view.open() as session: + adapter = LegacyMarkerQueryAdapter(session) df = ( - self._handle_backwards_compatibility(session.execute_lazy(query)) + session.execute_lazy(self._edgelist_query(adapter, components)) .collect() .to_pandas() ) @@ -87,32 +89,37 @@ def to_df(self) -> pd.DataFrame: def to_polars(self) -> pl.DataFrame: """Get the edgelist as a polars DataFrame.""" - query = self._query_builder.edgelist_query( - normalize_input_to_list(self.components) - ) + components = normalize_input_to_list(self.components) with self._view.open() as session: - df = self._handle_backwards_compatibility( - session.execute_lazy(query) + adapter = LegacyMarkerQueryAdapter(session) + df = session.execute_lazy( + self._edgelist_query(adapter, components) ).collect() return df def to_record_batches( self, batch_size: int = 1_000_000 ) -> Iterable[pa.RecordBatch]: - """Get the edgelist as a stream of pyarrow RecordBatches.""" - query = self._query_builder.edgelist_query( - normalize_input_to_list(self.components) - ) + """Get the edgelist as a stream of pyarrow RecordBatches. + + Legacy ``marker1`` and ``marker2`` columns are renamed to + ``marker_1`` and ``marker_2``, matching :meth:`to_polars`. + """ + components = normalize_input_to_list(self.components) with self._view.open() as session: - yield from session.execute_arrow_reader(query=query, batch_size=batch_size) + adapter = LegacyMarkerQueryAdapter(session) + yield from session.execute_arrow_reader( + query=self._edgelist_query(adapter, components), + batch_size=batch_size, + ) def _iterator(self) -> Iterable[tuple[str, pl.LazyFrame]]: with self._view.open() as session: + adapter = LegacyMarkerQueryAdapter(session) for component in self.components: - query = self._query_builder.edgelist_query([component]) yield ( component, - session.execute_lazy(query), + session.execute_lazy(self._edgelist_query(adapter, [component])), ) def iterator(self) -> Iterable[Component]: @@ -127,7 +134,7 @@ def iterator(self) -> Iterable[Component]: # here is that otherwise the object is not pickable, and thus not handled # well by the analysis manager. We should revisit this in the future. component_id=name, - frame=self._handle_backwards_compatibility(df).collect().lazy(), + frame=df.collect().lazy(), ) def __str__(self) -> str: diff --git a/src/pixelator/pna/pixeldataset/legacy/__init__.py b/src/pixelator/pna/pixeldataset/legacy/__init__.py new file mode 100644 index 000000000..9989b31c3 --- /dev/null +++ b/src/pixelator/pna/pixeldataset/legacy/__init__.py @@ -0,0 +1,8 @@ +"""Patches that make older ``.pxl`` files look like the current schema. + +Copyright © 2025 Pixelgen Technologies AB. +""" + +from pixelator.pna.pixeldataset.legacy.marker_columns import LegacyMarkerQueryAdapter + +__all__ = ["LegacyMarkerQueryAdapter"] diff --git a/src/pixelator/pna/pixeldataset/legacy/marker_columns.py b/src/pixelator/pna/pixeldataset/legacy/marker_columns.py new file mode 100644 index 000000000..b4cf61e22 --- /dev/null +++ b/src/pixelator/pna/pixeldataset/legacy/marker_columns.py @@ -0,0 +1,41 @@ +"""Adapt queries so an edgelist with legacy marker column names matches the current schema. + +Copyright © 2025 Pixelgen Technologies AB. +""" + +from __future__ import annotations + +from pixelator.pna.pixeldataset.io import PixelDataViewerSession, Query + +_LEGACY_MARKER_COLUMNS = (("marker1", "marker_1"), ("marker2", "marker_2")) + + +class LegacyMarkerQueryAdapter: + """Rename ``marker1`` and ``marker2`` on queries against a pre-rename edgelist. + + A current file is left unchanged. Construct one adapter per open session and + reuse it for every query in that session. + """ + + def __init__(self, session: PixelDataViewerSession) -> None: + """Record which legacy marker columns this session's edgelist still has.""" + columns = set( + session.execute_eager( + Query( + "SELECT column_name FROM (DESCRIBE SELECT * FROM edgelist)", + {}, + ) + )["column_name"].to_list() + ) + self._rename = ", ".join( + f"{old} AS {new}" for old, new in _LEGACY_MARKER_COLUMNS if old in columns + ) + + def adapt(self, query: Query) -> Query: + """Return ``query`` with legacy marker columns renamed, when any are present.""" + if not self._rename: + return query + return Query( + sql=(f"SELECT * RENAME ({self._rename}) FROM ({query.sql}) AS edgelist"), + params=query.params, + ) diff --git a/tests/pna/pixeldataset/test_anndata_helper.py b/tests/pna/pixeldataset/test_anndata_helper.py index e5daf6056..5b4b399a6 100644 --- a/tests/pna/pixeldataset/test_anndata_helper.py +++ b/tests/pna/pixeldataset/test_anndata_helper.py @@ -3,8 +3,10 @@ Copyright © 2025 Pixelgen Technologies AB. """ +import shutil from pathlib import Path +import duckdb import numpy as np import polars as pl import pytest @@ -331,3 +333,23 @@ def test_skips_bump_when_prerequisites_are_not_met( assert (adata_old[:, "MarkerC"].X == not_bumped[0][:, "MarkerC"].X).all() assert (adata_new[:, "MarkerC"].X == not_bumped[1][:, "MarkerC"].X).all() + + +def test_record_batches_rename_legacy_marker_columns(tmp_path: Path, pxl_file: Path): + """A legacy edgelist stream uses the same column names as ``to_polars``.""" + legacy = tmp_path / "legacy.pxl" + shutil.copy(pxl_file, legacy) + with duckdb.connect(str(legacy)) as connection: + connection.execute("ALTER TABLE edgelist RENAME COLUMN marker_1 TO marker1") + connection.execute("ALTER TABLE edgelist RENAME COLUMN marker_2 TO marker2") + dataset = PNAPixelDataset.from_files(legacy) + streamed = pl.concat( + [pl.from_arrow(batch) for batch in dataset.edgelist().to_record_batches()], + how="vertical", + ) + loaded = dataset.edgelist().to_polars() + assert "marker_1" in streamed.columns + assert "marker_2" in streamed.columns + assert "marker1" not in streamed.columns + assert "marker2" not in streamed.columns + assert streamed.sort(streamed.columns).equals(loaded.sort(loaded.columns))