Fix fallback representative assignment and streamline clustering

src/plinder/data/clusters.py

@@ -11,6 +11,7 @@
 from multiprocessing import get_context
 from pathlib import Path
 from shutil import copyfile, copytree, rmtree
+from tempfile import TemporaryDirectory
 from textwrap import dedent
 from time import time
 from typing import Any, Callable, Iterable, Literal, Mapping, Sequence, TypeVar, cast
@@ -23,6 +24,7 @@
 import networkit as nk
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 import pyarrow.parquet as pq
 from numpy.typing import NDArray
 
@@ -382,7 +384,9 @@ def reduce_threshold_component_edges(
     if not modes:
         raise ValueError("at least one component direction mode is required")
     nodes = sorted(set(map(str, all_nodes)))
-    node_indexes = {node: index for index, node in enumerate(nodes)}
+    node_indexes = pd.Series(
+        np.arange(len(nodes), dtype=np.int64), index=pd.Index(nodes)
+    )
     accumulators = {
         directed: {
             threshold: _ComponentEdgeAccumulator(
@@ -2833,6 +2837,8 @@ def select(
     best_similarity[representatives] = 100.0
 
     def assign(threshold: int, only_unassigned: bool) -> None:
+        # Freeze eligibility so later representatives can improve fallback hits.
+        eligible = best_representative < 0 if only_unassigned else ~selected
         for representative_index in representatives:
             representative = int(representative_index)
             start = int(index.offsets[representative])
@@ -2842,10 +2848,8 @@ def assign(threshold: int, only_unassigned: bool) -> None:
             keep = (
                 (queries != representative)
                 & (similarities >= threshold)
-                & ~selected[queries]
+                & eligible[queries]
             )
-            if only_unassigned:
-                keep &= best_representative[queries] < 0
             queries = queries[keep]
             similarities = similarities[keep]
             current_representatives = best_representative[queries]
@@ -3030,6 +3034,43 @@ def _directed_cover_labels(
     return labels
 
 
+def _write_target_sorted_edges(
+    connection: Any, query: str, staged_edges: Path, temporary_root: Path
+) -> None:
+    """Sort bounded target ranges instead of sorting the entire score graph."""
+    schema = pa.schema(
+        [
+            ("query_node", pa.uint32()),
+            ("target_node", pa.uint32()),
+            ("similarity", pa.uint8()),
+        ]
+    )
+    with TemporaryDirectory(dir=temporary_root, prefix="cover-ranges-") as directory:
+        connection.execute(
+            f"COPY (SELECT *, target_node // 65536 AS target_range FROM ({query})) "
+            "TO ? (FORMAT PARQUET, COMPRESSION ZSTD, PARTITION_BY(target_range), "
+            "ROW_GROUP_SIZE 2000000)",
+            [directory],
+        )
+        ranges = sorted(
+            Path(directory).glob("target_range=*"),
+            key=lambda path: int(path.name.split("=")[1]),
+        )
+        with pq.ParquetWriter(staged_edges, schema, compression="zstd") as writer:
+            for index, partition in enumerate(ranges, start=1):
+                batches = connection.execute(
+                    "SELECT query_node,target_node,similarity FROM read_parquet(?) "
+                    "ORDER BY target_node,query_node",
+                    [str(partition / "*.parquet")],
+                ).fetch_record_batch(2_000_000)
+                for batch in batches:
+                    writer.write_batch(batch)
+                rmtree(partition)
+                LOG.info(
+                    "directed cover sorted target ranges: %d/%d", index, len(ranges)
+                )
+
+
 def _stage_directed_cover_edges(
     *,
     data_dir: Path,
@@ -3103,7 +3144,6 @@ def _stage_directed_cover_edges(
             target_node,
             similarity
         FROM labeled
-        ORDER BY target_node, query_node
         """
     )
     LOG.info(
@@ -3121,11 +3161,7 @@ def _stage_directed_cover_edges(
     incoming_counts = np.zeros(len(labels), dtype=np.int64)
     stage_started = time()
     try:
-        escaped_staged_edges = staged_edges.as_posix().replace("'", "''")
-        connection.execute(
-            f"COPY ({query}) TO '{escaped_staged_edges}' "
-            "(FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 2000000)"
-        )
+        _write_target_sorted_edges(connection, query, staged_edges, temporary_root)
         previous_target = -1
         for batch_index, batch in enumerate(
             pq.ParquetFile(staged_edges).iter_batches(

tests/data/test_clusters.py

@@ -1007,6 +1007,35 @@ def test_greedy_set_cover_selects_the_largest_residual_gain():
     assert groups == [("broad-other", nodes)]
 
 
+@pytest.mark.parametrize("empty", [False, True])
+def test_cover_edge_sort_preserves_rows_across_target_ranges(tmp_path, empty):
+    import duckdb
+
+    from plinder.data.clusters import _write_target_sorted_edges
+
+    edges = pd.DataFrame(
+        {
+            "query_node": pd.Series([9, 2, 1, 7, 7, 3], dtype="uint32"),
+            "target_node": pd.Series(
+                [131072, 65536, 65535, 0, 0, 65536], dtype="uint32"
+            ),
+            "similarity": pd.Series([30, 100, 90, 50, 50, 70], dtype="uint8"),
+        }
+    )
+    if empty:
+        edges = edges.iloc[:0]
+    connection = duckdb.connect()
+    connection.register("edges", edges)
+    target = tmp_path / "sorted.parquet"
+    try:
+        _write_target_sorted_edges(connection, "SELECT * FROM edges", target, tmp_path)
+    finally:
+        connection.close()
+    expected = edges.sort_values(["target_node", "query_node"]).reset_index(drop=True)
+    pd.testing.assert_frame_equal(pd.read_parquet(target), expected)
+    assert not list(tmp_path.glob("cover-ranges-*"))
+
+
 def test_directed_cover_uses_query_to_centroid_scores_and_reassigns():
     from plinder.data.clusters import _greedy_directed_centroid_cover
 
@@ -1084,6 +1113,70 @@ def gains(threshold):
     )
 
 
+@pytest.mark.parametrize("fallback_scores", [(60, 80), (80, 60), (80, 80)])
+def test_indexed_cover_assigns_fallback_to_best_representative(
+    tmp_path, fallback_scores
+):
+    from plinder.data.clusters import (
+        _build_outgoing_edge_index,
+        _greedy_directed_set_cover,
+        _greedy_indexed_directed_set_cover,
+        _IncomingEdgeIndex,
+    )
+
+    nodes = ["a", "b", "strict-a", "strict-b", "relaxed"]
+    incoming = [
+        [(2, 90), (4, fallback_scores[0])],
+        [(3, 90), (4, fallback_scores[1])],
+        [],
+        [],
+        [],
+    ]
+    index = _IncomingEdgeIndex(
+        offsets=np.r_[0, np.cumsum([len(edges) for edges in incoming])],
+        queries=np.asarray(
+            [q for edges in incoming for q, _ in edges], dtype=np.uint32
+        ),
+        similarities=np.asarray(
+            [s for edges in incoming for _, s in edges], dtype=np.uint8
+        ),
+    )
+    graph = nk.Graph(len(nodes), weighted=True, directed=True)
+    for target, edges in enumerate(incoming):
+        for query, similarity in edges:
+            graph.addEdge(query, target, similarity / 100)
+    selections, representatives, similarities, thresholds = (
+        _greedy_indexed_directed_set_cover(
+            index,
+            nodes,
+            primary_threshold=90,
+            fallback_threshold=50,
+            primary_gains=np.asarray([2, 2, 1, 1, 1]),
+            fallback_gains=np.asarray([3, 3, 1, 1, 1]),
+            outgoing_index=_build_outgoing_edge_index(
+                index,
+                target_index_path=tmp_path / "targets.bin",
+                similarity_index_path=tmp_path / "scores.bin",
+            ),
+        )
+    )
+    expected_selections, expected_assignments = _greedy_directed_set_cover(
+        graph, nodes, primary_threshold=90, fallback_threshold=50
+    )
+    assert selections == expected_selections
+    np.testing.assert_array_equal(
+        representatives, [item.representative for item in expected_assignments]
+    )
+    np.testing.assert_allclose(
+        similarities, [item.similarity for item in expected_assignments]
+    )
+    np.testing.assert_array_equal(
+        thresholds, [item.threshold for item in expected_assignments]
+    )
+    assert similarities[4] == max(fallback_scores)
+    assert representatives[4] == (1 if fallback_scores[1] > fallback_scores[0] else 0)
+
+
 def test_indexed_cover_preserves_selected_centroids_in_perfect_ties(tmp_path):
     from plinder.data.clusters import (
         _build_outgoing_edge_index,