# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""OFFSET against keyset pagination, under writes.

1,000,000 rows in SQLite. The time to fetch page 10,000 with OFFSET and with
a keyset condition; then page 1, a new row inserted at the top, page 2: the
duplicates each method returns.
"""
import argparse, json, sqlite3, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=1_000_000, help="Rows.")
p.add_argument("--page", type=int, default=100, help="Rows per page.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
db = sqlite3.connect(":memory:")
db.execute("CREATE TABLE c(id INTEGER PRIMARY KEY, name TEXT)")
db.executemany("INSERT INTO c(name) VALUES (?)", ((f"n{i}",) for i in range(args.rows)))
last = args.rows - args.page
def timed(sql):
    t = time.perf_counter()
    for _ in range(20): db.execute(sql).fetchall()
    return (time.perf_counter() - t) / 20 * 1000
t_off = timed(f"SELECT id FROM c ORDER BY id DESC LIMIT {args.page} OFFSET {last}")
t_key = timed(f"SELECT id FROM c WHERE id < {args.page + 1} ORDER BY id DESC LIMIT {args.page}")
p1 = [r[0] for r in db.execute(f"SELECT id FROM c ORDER BY id DESC LIMIT {args.page} OFFSET 0")]
db.execute("INSERT INTO c(name) VALUES ('new')")
p2 = [r[0] for r in db.execute(f"SELECT id FROM c ORDER BY id DESC LIMIT {args.page} OFFSET {args.page}")]
k2 = [r[0] for r in db.execute(f"SELECT id FROM c WHERE id < {p1[-1]} ORDER BY id DESC LIMIT {args.page}")]
rows = [{"method": "OFFSET", "page_10000_ms": t_off, "duplicates": len(set(p1) & set(p2))},
        {"method": "keyset (WHERE id < last seen)", "page_10000_ms": t_key, "duplicates": len(set(p1) & set(k2))}]
for r in rows: print(f"{r['method']:<30} page 10,000 in {r['page_10000_ms']:7.3f} ms   duplicates after one insert: {r['duplicates']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"methods": rows, "summary": [{"speedup": t_off / t_key}]}}, f)