# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]", "numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Three self-joins against a CSR traversal.

A power-law follower graph of 1 million vertices (preferential attachment, 5
edges per new vertex) loaded into PostgreSQL as follows(follower_id,
followee_id) with an index. One- to four-hop chained self-joins from a median
vertex and from the top hub: time and frontier size. Then the same depths as a
level-synchronous BFS over CSR arrays in NumPy.
Needs the PostgreSQL service from the compose file (port 56002).
"""
import argparse, json, time
import numpy as np, psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:56002/postgres", help="Connection string.")
p.add_argument("--n", type=int, default=1_000_000, help="Vertices.")
p.add_argument("--m", type=int, default=5, help="Edges per new vertex.")
p.add_argument("--depth", type=int, default=4, help="Hops.")
p.add_argument("--timeout-s", type=int, default=120, help="Statement timeout for the joins.")
p.add_argument("--max-paths", type=float, default=5e7, help="Skip a SQL hop expected to enumerate more paths than this.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = np.random.default_rng(0)
targets = np.empty(args.n * args.m * 2, dtype=np.int64); src, dst = [], []; filled = 0
for v in range(args.m, args.n):                                          # preferential attachment: pick ends of existing edges
    picks = targets[rng.integers(0, filled, args.m)] if filled else np.arange(args.m)
    src.append(np.full(args.m, v)); dst.append(picks)
    targets[filled:filled + args.m] = picks; targets[filled + args.m:filled + 2 * args.m] = v; filled += 2 * args.m
s, d = np.concatenate(src), np.concatenate(dst)
s, d = np.concatenate([s, d]), np.concatenate([d, s])                   # both directions: follows and followed-by
order = np.argsort(s, kind="stable"); s, d = s[order], d[order]
off = np.zeros(args.n + 1, dtype=np.int64); np.add.at(off, s + 1, 1); off = np.cumsum(off); nbrs = d
deg = np.diff(off); hub = int(deg.argmax()); median_v = int(np.argsort(deg)[args.n // 2])
c = psycopg.connect(args.dsn, autocommit=True)
c.execute("DROP TABLE IF EXISTS follows"); c.execute("CREATE TABLE follows (follower_id int, followee_id int)")
with c.cursor().copy("COPY follows FROM STDIN") as cp:
    for a, b in zip(s.tolist(), d.tolist()): cp.write_row((a, b))
c.execute("CREATE INDEX ON follows (follower_id)"); c.execute("ANALYZE follows"); c.execute(f"SET statement_timeout = '{args.timeout_s}s'")
def sql_hops(start, k):
    joins = " ".join(f"JOIN follows f{i} ON f{i}.follower_id = f{i - 1}.followee_id" for i in range(1, k))
    q = f"SELECT count(*), count(DISTINCT f{k - 1}.followee_id) FROM follows f0 {joins} WHERE f0.follower_id = {start}"
    t = time.perf_counter()
    try: total, distinct = c.execute(q).fetchone(); return time.perf_counter() - t, total, distinct
    except psycopg.errors.QueryCanceled: return None, None, None
def bfs(start, k):
    visited = np.zeros(args.n, dtype=bool); visited[start] = True; frontier = np.array([start]); out = []
    t = time.perf_counter()
    for hop in range(k):
        nb = np.concatenate([nbrs[off[v]:off[v + 1]] for v in frontier.tolist()]) if len(frontier) else np.array([], dtype=np.int64)
        nb = np.unique(nb[~visited[nb]]); visited[nb] = True; frontier = nb
        out.append((hop + 1, time.perf_counter() - t, len(frontier), int(visited.sum())))
    return out
rows = []
for label, v in (("median vertex", median_v), ("top hub", hub)):
    paths = 1
    for hop, el, front, reached in bfs(v, args.depth):
        if paths * deg.mean() * 2 > args.max_paths: sj, total, distinct = None, None, None   # the join would not fit in memory
        else: sj, total, distinct = sql_hops(v, hop)
        paths = total if total else paths * deg.mean() * 2
        rows.append({"start": label, "hops": hop, "sql_s": sj, "sql_paths": total, "csr_s": el, "frontier": front, "reached": reached})
summary = {"vertices": args.n, "edges": int(len(s) // 2), "hub_degree": int(deg[hub]), "median_degree": int(deg[median_v]), "mean_degree": float(deg.mean())}
print(f"{args.n:,} vertices, {summary['edges']:,} edges; degree: median vertex {summary['median_degree']}, hub {summary['hub_degree']:,}")
for r in rows: print(f"{r['start']:<14} {r['hops']} hops: self-joins {('%.2f s' % r['sql_s']) if r['sql_s'] is not None else 'skipped':>10} ({r['sql_paths'] or 0:>12,} paths)   CSR BFS {r['csr_s']:6.2f} s, new frontier {r['frontier']:>9,}, reached {r['reached']:>9,}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"hops": rows}}, f)