"""
The queries are modeled after the shape an ORM or BI layer may produce: a
filter on a related table composed as a to-many join, with DISTINCT or
count(DISTINCT) to undo the fanout that join causes.  A person writing these
queries would probably use EXISTS, which is timed as well. That may represent a
realistic ceiling.

Each case is timed with enable_eager_aggregate off and then on so the
comparison is the patch against master's plan for the same query text.
Timing is wall clock around the statement, with the output discarded 
server-side.

Every case is checked for equivalence with the GUC off against on, and the
join form against its EXISTS form, before anything is timed.
"""

import os
import statistics
import sys
import time

import psycopg2

FANOUTS = [int(f) for f in os.environ.get("BENCH_FANOUTS", "1,2,8,32").split(",")]
REPS = int(os.environ.get("BENCH_REPS", "5"))
GUC = "enable_eager_aggregate"

# The default, so that a hash join over the fanned-out rows batches as it
# would out of the box.  Affects how much hash table fits in memory
WORK_MEM = os.environ.get("BENCH_WORK_MEM", "4MB")

NPOSTS = 100_000
NAUTHORS = 10_000

# The third hop squares the fanout, so reaction would reach 
# 100M+ rows at 32.  The shape is settled by 8.
DEEP_MAX_FANOUT = 8

# (label, join form, equivalent EXISTS form)
CASES = [
    ("already an EXISTS (control)",
     """SELECT p.id, p.title FROM post p
        WHERE EXISTS (SELECT 1 FROM comment c
                      WHERE c.post_id = p.id AND c.spam)""",
     None),

    ("DISTINCT over 1:many join",
     """SELECT DISTINCT p.id, p.author_id, p.title, p.published
        FROM post p JOIN comment c ON c.post_id = p.id
        WHERE c.spam""",
     """SELECT p.id, p.author_id, p.title, p.published
        FROM post p
        WHERE EXISTS (SELECT 1 FROM comment c
                      WHERE c.post_id = p.id AND c.spam)"""),

    ("count(DISTINCT) over 1:many join",
     """SELECT count(DISTINCT p.id)
        FROM post p JOIN comment c ON c.post_id = p.id
        WHERE c.spam""",
     """SELECT count(*) FROM post p
        WHERE EXISTS (SELECT 1 FROM comment c
                      WHERE c.post_id = p.id AND c.spam)"""),

    ("GROUP BY with max()",
     """SELECT p.id, max(p.title)
        FROM post p JOIN comment c ON c.post_id = p.id
        WHERE c.spam
        GROUP BY p.id""",
     """SELECT p.id, max(p.title) FROM post p
        WHERE EXISTS (SELECT 1 FROM comment c
                      WHERE c.post_id = p.id AND c.spam)
        GROUP BY p.id"""),

    ("DISTINCT, driver filtered",
     """SELECT DISTINCT p.id, p.title
        FROM post p JOIN comment c ON c.post_id = p.id
        WHERE c.spam AND p.published""",
     """SELECT p.id, p.title FROM post p
        WHERE p.published
          AND EXISTS (SELECT 1 FROM comment c
                      WHERE c.post_id = p.id AND c.spam)"""),

    # A bridge that fans out, joined on to a dimension.
    ("many to many",
     """SELECT DISTINCT p.id, p.title
        FROM post p
             JOIN comment c ON c.post_id = p.id
             JOIN author a2 ON a2.id = c.author_id
        WHERE a2.country LIKE 'c%'""",
     """SELECT p.id, p.title FROM post p
        WHERE EXISTS (SELECT 1
                      FROM comment c JOIN author a2 ON a2.id = c.author_id
                      WHERE c.post_id = p.id AND a2.country LIKE 'c%')"""),

    ("many to many, count(DISTINCT)",
     """SELECT count(DISTINCT p.id)
        FROM post p
             JOIN comment c ON c.post_id = p.id
             JOIN author a2 ON a2.id = c.author_id
        WHERE a2.country LIKE 'c%'""",
     """SELECT count(*) FROM post p
        WHERE EXISTS (SELECT 1
                      FROM comment c JOIN author a2 ON a2.id = c.author_id
                      WHERE c.post_id = p.id AND a2.country LIKE 'c%')"""),

    ("two-hop chain, DISTINCT",
     """SELECT DISTINCT a.id, a.name, a.country
        FROM author a
             JOIN post p ON p.author_id = a.id
             JOIN comment c ON c.post_id = p.id
        WHERE c.spam""",
     """SELECT a.id, a.name, a.country FROM author a
        WHERE EXISTS (SELECT 1
                      FROM post p JOIN comment c ON c.post_id = p.id
                      WHERE p.author_id = a.id AND c.spam)"""),

    ("two-hop chain, count(DISTINCT)",
     """SELECT count(DISTINCT a.id)
        FROM author a
             JOIN post p ON p.author_id = a.id
             JOIN comment c ON c.post_id = p.id
        WHERE c.spam""",
     """SELECT count(*) FROM author a
        WHERE EXISTS (SELECT 1
                      FROM post p JOIN comment c ON c.post_id = p.id
                      WHERE p.author_id = a.id AND c.spam)"""),

    # The same phrases one hop deeper
    ("three-hop chain, DISTINCT",
     """SELECT DISTINCT a.id, a.name, a.country
        FROM author a
             JOIN post p ON p.author_id = a.id
             JOIN comment c ON c.post_id = p.id
             JOIN reaction r ON r.comment_id = c.id
        WHERE r.flagged""",
     """SELECT a.id, a.name, a.country FROM author a
        WHERE EXISTS (SELECT 1
                      FROM post p
                           JOIN comment c ON c.post_id = p.id
                           JOIN reaction r ON r.comment_id = c.id
                      WHERE p.author_id = a.id AND r.flagged)"""),

    ("three-hop chain, count(DISTINCT)",
     """SELECT count(DISTINCT a.id)
        FROM author a
             JOIN post p ON p.author_id = a.id
             JOIN comment c ON c.post_id = p.id
             JOIN reaction r ON r.comment_id = c.id
        WHERE r.flagged""",
     """SELECT count(*) FROM author a
        WHERE EXISTS (SELECT 1
                      FROM post p
                           JOIN comment c ON c.post_id = p.id
                           JOIN reaction r ON r.comment_id = c.id
                      WHERE p.author_id = a.id AND r.flagged)"""),
]

USE_COPY = True


def needs_reaction(label):
    return label.startswith("three-hop")


def connect(dbname):
    last = None
    for host in (os.environ.get("PGHOST"), "/tmp", "/var/run/postgresql"):
        if host is None:
            continue
        try:
            conn = psycopg2.connect(dbname=dbname, host=host,
                                    user=os.environ.get("PGUSER"))
            conn.autocommit = True
            return conn
        except psycopg2.Error as exc:
            last = exc
    raise SystemExit("could not connect: %s" % last)


def create_database():
    conn = connect("postgres")
    cur = conn.cursor()
    cur.execute("SELECT 1 FROM pg_database WHERE datname = 'fanout'")
    if cur.fetchone() is None:
        cur.execute("CREATE DATABASE fanout")
    conn.close()


def load_base(cur):
    cur.execute("DROP TABLE IF EXISTS reaction, comment, post, author")
    cur.execute("""
        CREATE TABLE author (
            id integer PRIMARY KEY,
            name text NOT NULL,
            country text NOT NULL)
    """)
    cur.execute("""
        CREATE TABLE post (
            id integer PRIMARY KEY,
            author_id integer NOT NULL REFERENCES author,
            title text NOT NULL,
            published boolean NOT NULL)
    """)

    cur.execute("""
        INSERT INTO author
        SELECT g, 'author' || g, 'c' || (g %% 20)
        FROM generate_series(1, %s) g
    """, (NAUTHORS,))
    # Ten posts per author, fixed by the corpus
    cur.execute("""
        INSERT INTO post
        SELECT g, 1 + (g %% %s), 'title' || g, (g %% 3) = 0
        FROM generate_series(1, %s) g
    """, (NAUTHORS, NPOSTS))

    cur.execute("CREATE INDEX ON post (author_id)")
    cur.execute("ANALYZE author, post")


def load_comments(cur, fanout):
    """Give every post `fanout` matching comments."""
    cur.execute("DROP TABLE IF EXISTS reaction")
    cur.execute("DROP TABLE IF EXISTS comment")
    cur.execute("""
        CREATE TABLE comment (
            id serial PRIMARY KEY,
            post_id integer NOT NULL REFERENCES post,
            author_id integer NOT NULL REFERENCES author,
            body text NOT NULL,
            spam boolean NOT NULL)
    """)
    cur.execute("""
        INSERT INTO comment (post_id, author_id, body, spam)
        SELECT p.id, 1 + ((p.id + g) %% %s), 'body', true
        FROM post p, generate_series(1, %s) g
    """, (NAUTHORS, fanout))
    cur.execute("CREATE INDEX ON comment (post_id)")
    cur.execute("ANALYZE comment")


def load_reactions(cur, fanout):
    """Give every comment `fanout` matching reactions: the chain's third hop."""
    cur.execute("""
        CREATE TABLE reaction (
            id serial PRIMARY KEY,
            comment_id integer NOT NULL REFERENCES comment,
            author_id integer NOT NULL REFERENCES author,
            kind text NOT NULL,
            flagged boolean NOT NULL)
    """)
    cur.execute("""
        INSERT INTO reaction (comment_id, author_id, kind, flagged)
        SELECT c.id, 1 + ((c.id + g) %% %s), 'kind', true
        FROM comment c, generate_series(1, %s) g
    """, (NAUTHORS, fanout))
    cur.execute("CREATE INDEX ON reaction (comment_id)")
    cur.execute("ANALYZE reaction")


def pick_timing_method(cur):
    """Prefer discarding output server side; fall back to fetching it."""
    global USE_COPY
    try:
        cur.execute("COPY (SELECT 1) TO PROGRAM 'cat > /dev/null'")
    except psycopg2.Error:
        USE_COPY = False
        print("COPY TO PROGRAM unavailable; timing includes fetching rows")


def run_once(cur, sql):
    """Wall clock milliseconds for one execution, output discarded."""
    start = time.perf_counter()
    if USE_COPY:
        cur.execute("COPY (%s) TO PROGRAM 'cat > /dev/null'" % sql)
    else:
        cur.execute(sql)
        cur.fetchall()
    return (time.perf_counter() - start) * 1000.0


def measure(cur, sql, exists_sql):
    """Time the GUC off, on, and the hand-written form, interleaved."""
    for setting in ("off", "on"):
        cur.execute("SET %s = %s" % (GUC, setting))
        run_once(cur, sql)
    if exists_sql:
        run_once(cur, exists_sql)

    off_ms, on_ms, ex_ms = [], [], []
    for _ in range(REPS):
        cur.execute("SET %s = off" % GUC)
        off_ms.append(run_once(cur, sql))
        cur.execute("SET %s = on" % GUC)
        on_ms.append(run_once(cur, sql))
        if exists_sql:
            cur.execute("SET %s = off" % GUC)
            ex_ms.append(run_once(cur, exists_sql))

    return (statistics.median(off_ms), statistics.median(on_ms),
            statistics.median(ex_ms) if ex_ms else None)


def plan_of(cur, sql, setting):
    cur.execute("SET %s = %s" % (GUC, setting))
    cur.execute("EXPLAIN (COSTS OFF) " + sql)
    return "\n".join(r[0] for r in cur.fetchall())


def engagement(cur, sql):
    """Did the GUC change the plan, and by pushing down what?"""
    off = plan_of(cur, sql, "off")
    on = plan_of(cur, sql, "on")
    if off == on:
        return "-"
    if "Semi Join" in on and "Semi Join" not in off:
        return "semi"
    return "fold"


def result_hash(cur, sql):
    cur.execute("SELECT md5(coalesce(string_agg(t::text, '|' ORDER BY "
                "t::text), 'empty')) FROM (%s) t" % sql)
    return cur.fetchone()[0]


def equivalent(cur, join_sql, semi_sql):
    cur.execute("""
        SELECT (SELECT count(*) FROM ((%s) EXCEPT (%s)) d)
             + (SELECT count(*) FROM ((%s) EXCEPT (%s)) d)
    """ % (join_sql, semi_sql, semi_sql, join_sql))
    return cur.fetchone()[0] == 0


def print_table(title, rows):
    """One column per fanout, however many were run."""
    row_fmt = " %-32s" + " %8s" * len(FANOUTS)

    print("\n=== %s ===\n" % title)
    print(row_fmt % (("query          fanout",)
                     + tuple("=%d" % f for f in FANOUTS)))
    print(" " + "-" * (33 + 9 * len(FANOUTS)))
    for label, cells in rows:
        print(row_fmt % ((label,) + tuple(cells)))


def summarize(results):
    rows = []
    for label, _, _ in CASES:
        cells = []
        for fanout in FANOUTS:
            got = results.get((label, fanout))
            if got is None:
                cells.append("DNR")
                continue
            speedup, _, pushed = got
            cells.append("%.2fx%s" % (speedup, "*" if pushed == "-" else ""))
        rows.append((label, cells))
    print()
    print_table("Delivered speedup vs master", rows)
    print("\n  * the plan did not change | DNR -> not run")

    rows = []
    for label, _, _ in CASES:
        rows.append((label, [results.get((label, f), (0, 0, "DNR"))[2]
                             for f in FANOUTS]))
    print_table("Which pushdown the planner chose", rows)

    rows = []
    for label, _, exists_sql in CASES:
        if exists_sql is None:
            continue
        cells = []
        for fanout in FANOUTS:
            got = results.get((label, fanout))
            cells.append("%.2fx" % got[1] if got and got[1] else "DNR")
        rows.append((label, cells))
    print_table("Hand-written EXISTS vs master: the ceiling", rows)

    rows = []
    for label, _, exists_sql in CASES:
        if exists_sql is None:
            continue
        cells = []
        for fanout in FANOUTS:
            got = results.get((label, fanout))
            if got is None or not got[1]:
                cells.append("DNR")
                continue
            speedup, ceiling, _ = got
            cells.append("%d%%" % round(100.0 * speedup / ceiling))
        rows.append((label, cells))
    print_table("Share of that ceiling the pushdown captures", rows)


def main():
    create_database()
    conn = connect("fanout")
    cur = conn.cursor()

    cur.execute("SET max_parallel_workers_per_gather = 0")
    cur.execute("SET work_mem = %s", (WORK_MEM,))
    pick_timing_method(cur)

    print("loading base tables (%d posts, %d authors, ten posts per author)..."
          % (NPOSTS, NAUTHORS), flush=True)
    load_base(cur)

    results = {}
    checked = set()

    for fanout in FANOUTS:
        deep = fanout <= DEEP_MAX_FANOUT

        print("\nloading %d comments per post..." % fanout, flush=True)
        load_comments(cur, fanout)
        if deep:
            print("loading %d reactions per comment (%d rows)..."
                  % (fanout, NPOSTS * fanout * fanout), flush=True)
            load_reactions(cur, fanout)
        else:
            print("skipping reaction and the three-hop cases above fanout %d"
                  % DEEP_MAX_FANOUT, flush=True)

        print("\n=== fanout %d, serial, work_mem %s, median of %d ==="
              % (fanout, WORK_MEM, REPS))
        head = ("%-34s %9s %9s %8s %9s %8s %6s" %
                ("query", "off ms", "on ms", "speedup", "EXISTS",
                 "ceiling", "pushed"))
        print(head)
        print("-" * len(head))

        for label, sql, exists_sql in CASES:
            if not deep and needs_reaction(label):
                continue

            cur.execute("SET %s = off" % GUC)
            if exists_sql and label not in checked:
                if not equivalent(cur, sql, exists_sql):
                    print("MISMATCH: %s differs from its EXISTS form" % label)
                    return 1
                checked.add(label)

            off_hash = result_hash(cur, sql)
            cur.execute("SET %s = on" % GUC)
            on_hash = result_hash(cur, sql)
            if off_hash != on_hash:
                print("MISMATCH: %s changes result with the GUC on" % label)
                return 1

            pushed = engagement(cur, sql)
            off_ms, on_ms, ex_ms = measure(cur, sql, exists_sql)
            speedup = off_ms / on_ms if on_ms else 0
            ceiling = off_ms / ex_ms if ex_ms else None

            print("%-34s %9.1f %9.1f %7.2fx %9s %8s %6s" % (
                label, off_ms, on_ms, speedup,
                "%.1f" % ex_ms if ex_ms else "-",
                "%.2fx" % ceiling if ceiling else "-",
                pushed), flush=True)

            results[(label, fanout)] = (speedup, ceiling, pushed)

    summarize(results)
    conn.close()
    return 0


if __name__ == "__main__":
    sys.exit(main())