Top-N heap against full sort

# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Top-N: heap against full sort, and a LIMIT that backfires.

Part 1: the 100 largest of 20 million random numbers with sorted(data)[-100:]
and with heapq.nlargest(100, data), each in its own process so that its peak
memory can be read. Part 2 (PostgreSQL): a table where created_at and id grow
together and all 'rare' rows are at the end; WHERE flag = 'rare' ORDER BY id
LIMIT 10, with and without the LIMIT.
Needs the PostgreSQL service from the compose file (port 55707).
"""
import argparse, json, re, resource, subprocess, sys
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55707/postgres", help="Connection string.")
p.add_argument("--n", type=int, default=20_000_000, help="Random numbers.")
p.add_argument("--rows", type=int, default=5_000_000, help="Rows in the PostgreSQL table.")
p.add_argument("--rare", type=int, default=5_000, help="Rare rows, all at the end of the table.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

CHILD = r"""
import heapq, random, resource, sys, time
random.seed(0); data = [random.random() for _ in range(int(sys.argv[2]))]
base = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
t = time.perf_counter()
top = sorted(data)[-100:] if sys.argv[1] == "sort" else heapq.nlargest(100, data)
el = time.perf_counter() - t
peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
unit = 2**20 if sys.platform == "darwin" else 2**10
print(el, (peak - base) / unit)
"""
methods = []
for m, label in (("sort", "sorted(data)[-100:]"), ("heap", "heapq.nlargest(100, data)")):
    el, extra = map(float, subprocess.run([sys.executable, "-c", CHILD, m, str(args.n)], capture_output=True, text=True, check=True).stdout.split())
    methods.append({"method": label, "seconds": el, "extra_peak_mb": extra})

c = psycopg.connect(args.dsn, autocommit=True)
c.execute("DROP TABLE IF EXISTS events")
c.execute(f"""CREATE TABLE events AS SELECT g AS id, TIMESTAMP '2025-01-01' + g * INTERVAL '1 second' AS created_at,
              CASE WHEN g > {args.rows} - {args.rare} THEN 'rare' ELSE 'common' END AS flag FROM generate_series(1, {args.rows}) g""")
c.execute("ALTER TABLE events ADD PRIMARY KEY (id)"); c.execute("ANALYZE events")
plans = []
for label, q in (("with LIMIT 10", "SELECT * FROM events WHERE flag = 'rare' ORDER BY id LIMIT 10"),
                 ("without LIMIT", "SELECT * FROM events WHERE flag = 'rare' ORDER BY id")):
    txt = "\n".join(r[0] for r in c.execute("EXPLAIN (ANALYZE) " + q).fetchall())
    top = txt.splitlines()[0]; node = "index scan in id order" if "Index Scan" in txt else "scan + sort"
    plans.append({"query": label, "plan": node, "ms": float(re.search(r"Execution Time: ([\d.]+)", txt).group(1))})
for r in methods: print(f"{r['method']:<28} {r['seconds']:6.2f} s   extra peak memory {r['extra_peak_mb']:7.1f} MB")
for r in plans: print(f"rare rows at the end, {r['query']:<14}: {r['plan']:<24} {r['ms']:9.1f} ms")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"methods": methods, "plans": plans}}, f)

# PostgreSQL for the top-n-heap-against-full-sort lab.
services:
  postgres:
    image: postgres:17.6
    command: ["postgres", "-c", "shared_buffers=256MB", "-c", "max_parallel_workers_per_gather=0"]
    environment:
      POSTGRES_PASSWORD: lab
    ports: ["55707:5432"]
    healthcheck:
      test: ["CMD", "pg_isready", "-U", "postgres"]
      interval: 1s
      retries: 30