fix(edgelist): rename legacy marker columns in record batches

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.

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,45 +77,49 @@ 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()
             )
         return df
 
     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:

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"]

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,
+        )

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))