bench_fetchall.py

from time import sleep, monotonic
import statistics as st
import psycopg
import psycopg2

DSN = 'dbname=psycopg_test host=localhost user=postgres password=mysecretpassword'

table = """
CREATE TEMPORARY TABLE int_testtable (
    id serial PRIMARY KEY,
    i1 integer,
    i2 integer,
    i3 integer,
    i4 integer,
    i5 integer,
    i6 integer,
    i7 integer,
    i8 integer,
    i9 integer,
    i10 integer
)
"""

drop = """DROP TABLE IF EXISTS int_testtable"""

insert = """
INSERT INTO int_testtable (i1,i2,i3,i4,i5,i6,i7,i8,i9,i10)
SELECT
    1, 10, 100, 1000, 10000, 100000, 1000000, 10000000, 100000000, 1000000000
FROM generate_series(1, 1000000);
"""

select_i1 = """SELECT i1 FROM int_testtable LIMIT %(limit)s"""
select_i10 = """SELECT i10 FROM int_testtable LIMIT %(limit)s"""
select_all = """SELECT * FROM int_testtable LIMIT %(limit)s"""


def setup_testtable(conn):
    with conn.cursor() as c:
        c.execute(drop)
        c.execute(table)
        c.execute(insert)


def measure(func, args=None, kwargs=None, msg=None, repeat=10, relax=1.0):
    if repeat < 5:
        raise Exception('repeat must be greater than 4')
    durations = []
    for _ in range(repeat):
        start = monotonic()
        func(*(args or []), **(kwargs or {}))
        end = monotonic()
        durations.append(end - start)
        sleep(relax)
    mean = st.mean(durations)
    median = st.median(durations)
    sdev = st.stdev(durations)
    cv = sdev / mean
    msg = f"{func.__name__} {msg}" if msg else func.__name__
    print(f"{msg}: {round(mean, 5)} ±{round(sdev, 5)} (cv: {round(cv, 3)})")
    return {
        'raw': durations,
        'mean': mean,
        'median': median,
        'sdev': sdev,
        'cv': sdev / mean
    }


def bench_psycopg2():
    with psycopg2.connect(DSN) as conn:
        setup_testtable(conn)
        conn.commit()
        sleep(5)

        def v2_select_all_1M():
            with conn.cursor() as c:
                c.execute(select_all, {'limit': 1000000})
                data = c.fetchall()
        measure(v2_select_all_1M, repeat=5, relax=1.0)

        def v2_select_i1_1000():
            with conn.cursor() as c:
                c.execute(select_i1, {'limit': 1000})
                data = c.fetchall()
        measure(v2_select_i1_1000, repeat=50, relax=.1)


bench_psycopg2()


def bench_psycopg():
    with psycopg.connect(DSN) as conn:
        setup_testtable(conn)
        conn.commit()
        sleep(5)

        def v3_select_all_1M():
            with conn.cursor() as c:
                c.execute(select_all, {'limit': 1000000})
                data = c.fetchall()
        measure(v3_select_all_1M, repeat=5, relax=1.0)

        def v3_select_i1_10():
            with conn.cursor() as c:
                c.execute(select_i1, {'limit': 10})
                data = c.fetchall()
        measure(v3_select_i1_10, repeat=50, relax=.1)

        def v3_select_all_1M_best(page_size):
            with conn.cursor() as c:
                c.execute(select_all, {'limit': 1000000})
                c._tx._page_size = page_size
                data = c.fetchall()
        for page_size in [10, 20, 40, 80, 160, 320, 640, 1280]:
            measure(v3_select_all_1M_best, args=(page_size,),  msg=str(page_size), repeat=20, relax=1)
        
        def v3_select_i1_1000_best(page_size):
            with conn.cursor() as c:
                c.execute(select_i1, {'limit': 1000})
                c._tx._page_size = page_size
                data = c.fetchall()
        for page_size in [10, 20, 40, 80, 160, 320, 640, 1280]:
            measure(v3_select_i1_1000_best, args=(page_size,),  msg=str(page_size), repeat=20, relax=.1)
        
        def v3_select_i10_1000_best(page_size):
            with conn.cursor() as c:
                c.execute(select_i10, {'limit': 1000})
                c._tx._page_size = page_size
                data = c.fetchall()
        for page_size in [10, 20, 40, 80, 160, 320, 640, 1280]:
            measure(v3_select_i10_1000_best, args=(page_size,),  msg=str(page_size), repeat=20, relax=.1)

bench_psycopg()