From 148a72210a0fdb92c00e74081fe11f0a56847e63 Mon Sep 17 00:00:00 2001
From: ptajvar <ptajvar@gmail.com>
Date: Tue, 4 Aug 2026 16:20:26 +0200
Subject: [PATCH] ensure matching size schemas in
 filter_connected_components_by_size

---
 CHANGELOG.md                                  |  1 +
 .../pna/graph/component_recovery_utils.py     |  5 ++-
 .../graph/test_component_recovery_utils.py    | 38 +++++++++++++++++++
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 73083a3b2..ad442fdaa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -55,6 +55,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `--multiplet-recovery` and a `--component-size-max-threshold` that discards large components.
 - `write_parameters_file` now includes Click positional arguments (e.g. multi-file inputs) under
   `cli.arguments` in `*.meta.json`, not only options.
+- `filter_connected_components_by_size` now ensures matching size schemas between discarded and remaining components.
 
 ## [0.29.0] - 2026-06-15
 
diff --git a/src/pixelator/pna/graph/component_recovery_utils.py b/src/pixelator/pna/graph/component_recovery_utils.py
index bd4e4e3e2..1577e646a 100644
--- a/src/pixelator/pna/graph/component_recovery_utils.py
+++ b/src/pixelator/pna/graph/component_recovery_utils.py
@@ -605,7 +605,10 @@ def filter_connected_components_by_size(
         ConnectedComponentException: If no components remain after filtering.
     """
     component_sizes = create_component_size_data_frame(input_edgelist_path)
-    component_sizes = pl.concat([component_sizes, discard_sizes], how="vertical")
+    component_sizes = pl.concat(
+        [component_sizes, discard_sizes.select(["component", "n_umi"])],
+        how="vertical",
+    )
 
     unique, counts = np.unique(
         component_sizes["n_umi"].cast(pl.Int32), return_counts=True
diff --git a/tests/pna/graph/test_component_recovery_utils.py b/tests/pna/graph/test_component_recovery_utils.py
index e337952c3..351973481 100644
--- a/tests/pna/graph/test_component_recovery_utils.py
+++ b/tests/pna/graph/test_component_recovery_utils.py
@@ -252,6 +252,44 @@ def test_filter_connected_components_by_size_includes_early_discards_in_pre_filt
     assert fraction_of_discarded_components == pytest.approx(1 - 2 / 4)
 
 
+def test_filter_connected_components_by_size_accepts_discard_sizes_with_n_edges(
+    tmp_path: Path,
+) -> None:
+    """discard_sizes as returned by write_hive_partitioned_edgelist_without_out_of_size_bound_components carries an n_edges column; it must not break the concat with the (component, n_umi) sizes frame."""
+    input_path = tmp_path / "component_filter_input.parquet"
+    pl.DataFrame(
+        {
+            "component": ["a", "a", "b", "c", "c"],
+            "umi1": ["u1", "u3", "v1", "w1", "w1"],
+            "umi2": ["u2", "u4", "v2", "w2", "w3"],
+        }
+    ).write_parquet(input_path)
+    discard_sizes = pl.DataFrame(
+        {
+            "component": ["huge"],
+            "n_umi": pl.Series([10], dtype=pl.UInt32),
+            "n_edges": pl.Series([20], dtype=pl.UInt32),
+        }
+    )
+    component_stats = GraphStatistics()
+
+    filtered_edgelist_path, stats = filter_connected_components_by_size(
+        input_edgelist_path=input_path,
+        component_size_threshold=(3, 4),
+        discard_sizes=discard_sizes,
+        component_stats=component_stats,
+        working_dir=tmp_path,
+    )
+
+    filtered = pl.scan_parquet(
+        filtered_edgelist_path, hive_schema={"component": pl.String}
+    ).collect()
+
+    assert set(filtered["component"].unique().to_list()) == {"a", "c"}
+    assert stats.component_count_pre_component_size_filtering == 4
+    assert stats.pre_filtering_component_sizes == {2: 1, 3: 1, 4: 1, 10: 1}
+
+
 def test_create_component_size_data_frame_computes_sizes_per_component(
     tmp_path: Path,
 ) -> None: