add graph builder

build_graph.py

···
+from collections import UserString
+import logging
+from typing import Dict, Optional, Set
+
+import click
+
+from config import CONFIG
+from indexer import FollowIndexer
+import indexer
+
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+
+logger = logging.getLogger(__name__)
+
+
[e-mail]
[e-mail](
+    "--ch-host",
+)
[e-mail](
+    "--ch-port",
+    type=int,
+)
[e-mail](
+    "--ch-user",
+)
[e-mail](
+    "--ch-pass",
+)
+def main(
+    ch_host: Optional[str],
+    ch_port: Optional[int],
+    ch_user: Optional[str],
+    ch_pass: Optional[str],
+):
+    logger.info("Building follow graph...")
+
+    indexer = FollowIndexer(
+        clickhouse_host=ch_host or CONFIG.clickhouse_host,
+        clickhouse_port=ch_port or CONFIG.clickhouse_port,
+        clickhouse_user=ch_user or CONFIG.clickhouse_user,
+        clickhouse_pass=ch_pass or CONFIG.clickhouse_pass,
+        batch_size=1000,
+    )
+
+    graph: Dict[str, Set[str]] = {}
+
+    def build_graph(did: str, subject: str):
+        if did not in graph:
+            graph[did] = set()
+
+        graph[did].add(subject)
+
+    indexer.stream_follows(build_graph)
+
+    prox_map = {}
+
+    for did in graph:
+        first = graph.get(did, set())
+
+        second: Set[str] = set()
+        for subject in first:
+            second.update(graph.get(subject, set()))
+
+        prox_map[did] = {
+            "hop1": first,
+            "hop2": second - first - {did},
+        }
+
+    import pickle
+
+    with open("prox_map.pkl", "wb") as f:
+        pickle.dump(prox_map, f)
+
+    logger.info(
+        f"Finished building proximity map, saved to prox_map.pkl. {len(prox_map):,} users in map."
+    )
+
+
+if __name__ == "__main__":
+    main()

indexer.py

···
 from datetime import datetime
 from threading import Lock
 from time import time
-from typing import Any, List, Optional
+from typing import Any, Callable, List, Optional
 
 import click
 from aiokafka import AIOKafkaConsumer, ConsumerRecord
···
                 batch_to_flush = self._unfollow_batch.copy()
                 self._unfollow_batch = []
                 self._flush_unfollows(batch_to_flush)
+
+    def stream_follows(self, cb: Callable[[str, str], None], batch_size: int = 100_000):
+        query = """
+        SELECT f.did, f.subject
+        FROM follows f
+        LEFT ANTI JOIN unfollows u ON f.uri = u.uri
+        """
+
+        try:
+            with self.client.query_row_block_stream(
+                query, settings={"max_block_size": batch_size}
+            ) as stream:
+                total_handled = 0
+                for block in stream:
+                    for row in block:
+                        cb(row[0], row[1])
+                        total_handled += 1
+
+                    if total_handled % 1_000_000 == 0:
+                        logger.info(f"Handled {total_handled:,} follows so far")
+                logger.info(f"Finished streaming {total_handled:,} follows")
+        except Exception as e:
+            logger.error(f"Error streaming follows: {e}")
 
 
 class Consumer: