test: refactor out get_clusters() into tests.util for external use

Author: NickCrews

mismo/cluster/test/test_connected_components.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
-from typing import Any
-
 import ibis
 from ibis.expr import types as ir
 import pandas as pd
 import pytest
 
 from mismo.cluster import connected_components
+from mismo.tests.util import get_clusters
 
 
 @pytest.fixture(params=["component", "cluster"])
@@ -145,14 +144,5 @@ def test_cc_max_iterations(table_factory):
 
 def _labels_to_clusters(
     labels: ir.Table, label_as: str = "component"
-) -> set[frozenset[Any]]:
-    labels = labels.rename(component=label_as)
-    assert labels.component.type() == ibis.dtype("int64")
-    df = labels.to_pandas()
-    cid_to_rid = {c: set() for c in set(df.component)}
-    for row in df.itertuples():
-        record_id = row.record_id
-        if isinstance(record_id, dict):
-            record_id = tuple(record_id.values())
-        cid_to_rid[row.component].add(record_id)
-    return {frozenset(records) for records in cid_to_rid.values()}
+) -> set[frozenset[int]]:
+    return get_clusters(labels[label_as], label=labels.record_id)

mismo/tests/util.py

@@ -2,6 +2,7 @@
 
 from typing import Literal
 
+import ibis
 from ibis.expr import types as ir
 import pandas as pd
 import pytest
@@ -50,3 +51,30 @@ def make_float_comparable(val):
     if isinstance(val, float) and not pd.isna(val):
         return pytest.approx(val, rel=1e-4)
     return val
+
+
+def get_clusters(
+    cluster_id: ibis.Column, *, label: ibis.Column | None = None
+) -> set[frozenset]:
+    """Convert a label column into a set of clusters.
+
+    Say you have a table of records, and one of the columns
+    is an ID that groups records together.
+    This function will return a set of frozensets, where each
+    frozenset is a cluster of record IDs.
+
+    You can either provide a label column to act as the record IDs,
+    or if not given, it will use `ibis.row_number()`.
+    """
+    if label is None:
+        label = ibis.row_number()
+    clusters = label.collect().over(group_by=cluster_id)
+
+    def make_hashable(cluster):
+        for record_id in cluster:
+            if isinstance(record_id, dict):
+                yield tuple(record_id.values())
+            else:
+                yield record_id
+
+    return {frozenset(make_hashable(c)) for c in clusters.execute()}