run_regression_test.py
text/x-python-script
#!/usr/bin/env python3
"""
Run prefetch regression tests.
Usage:
python run_regression_test.py [options]
Examples:
python run_regression_test.py --iterations=5
python run_regression_test.py --evict=off,pg --workers=0,2
python run_regression_test.py --columns=sequential,random --reset
"""
import argparse
import subprocess
import itertools
import random
import os
import psycopg
import pandas as pd
from math import gcd
def parse_args():
p = argparse.ArgumentParser(description='Run prefetch regression tests')
p.add_argument('--iterations', '-n', type=int, default=10,
help='Number of test iterations (default: 10)')
p.add_argument('--columns', '-c', type=str, default='sequential,periodic,random',
help='Columns to test (default: sequential,periodic,random)')
p.add_argument('--workers', '-w', type=str, default='0,2',
help='Worker counts to test (default: 0,2)')
p.add_argument('--evict', '-e', type=str, default='off',
help='Evict modes: off,pg,os or "all" (default: off)')
p.add_argument('--rows', '-r', type=int, default=100000,
help='Number of rows in test table (default: 100000)')
p.add_argument('--reset', action='store_true',
help='Reset tables before running')
p.add_argument('--dbname', '--db', '-d', type=str, default='postgres',
help='Database name (default: postgres)')
p.add_argument('--host', '-H', type=str, default='/tmp',
help='Database host (default: /tmp)')
p.add_argument('--port', '-p', type=int, default=None,
help='Database port (default: use socket)')
return p.parse_args()
def setup_tables(cur, num_rows, reset=False):
"""Create tables and populate test data."""
if reset:
cur.execute('DROP TABLE IF EXISTS prefetch_test_results')
cur.execute('DROP TABLE IF EXISTS prefetch_test_data')
cur.execute('''
CREATE EXTENSION IF NOT EXISTS pg_buffercache
''');
cur.execute('''
CREATE TABLE IF NOT EXISTS prefetch_test_results (
id SERIAL PRIMARY KEY,
run_timestamp TIMESTAMPTZ DEFAULT now(),
io_method TEXT NOT NULL,
num_workers INT NOT NULL DEFAULT 0,
prefetch_enabled BOOLEAN NOT NULL,
evict_mode TEXT NOT NULL DEFAULT 'off',
column_name TEXT NOT NULL,
iteration INT,
execution_time_ms NUMERIC,
rows_returned BIGINT,
blks_hit BIGINT,
blks_read BIGINT
)
''')
cur.execute('''
CREATE TABLE IF NOT EXISTS prefetch_test_data (
id SERIAL PRIMARY KEY,
sequential INT,
periodic INT,
random INT,
payload TEXT
)
''')
cur.execute("SELECT count(*) FROM prefetch_test_data")
row_count = cur.fetchone()[0]
if row_count == 0:
print(f"Populating test data ({num_rows} rows)...")
r = min(10000, num_rows // 5)
while gcd(r, num_rows) != 1:
r += 1
cur.execute(f'''
INSERT INTO prefetch_test_data (sequential, periodic, random, payload)
SELECT
i,
((i * {r}::bigint) % {num_rows} + 1)::int,
row_number() OVER (ORDER BY random()),
repeat('x', 200)
FROM generate_series(1, {num_rows}) i
ORDER BY i;
''')
cur.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_sequential ON prefetch_test_data(sequential)")
cur.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_periodic ON prefetch_test_data(periodic)")
cur.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_random ON prefetch_test_data(random)")
print("Data populated.")
else:
print(f"Test data exists: {row_count} rows")
def purge_os_cache():
"""Purge OS filesystem cache using external script."""
script = os.path.join(os.path.dirname(__file__), 'purge_cache.sh')
try:
subprocess.run([script], capture_output=True, timeout=10)
except:
pass
def evict_pg_buffers(cur):
"""Evict PostgreSQL shared buffers for test table."""
cur.execute('''
SELECT count(pg_buffercache_evict(bufferid))
FROM pg_buffercache b
JOIN pg_class c ON b.relfilenode = pg_relation_filenode(c.oid)
WHERE c.relname = 'prefetch_test_data'
''').fetchall()
def apply_eviction(cur, evict_mode):
"""Apply eviction based on mode."""
if evict_mode == 'pg':
evict_pg_buffers(cur)
elif evict_mode == 'os':
evict_pg_buffers(cur)
purge_os_cache()
def run_test(cur, column_name, prefetch_enabled, num_workers, evict_mode, iteration):
"""Run a single test and record results."""
cur.execute(f"SET enable_indexscan_prefetch = {'on' if prefetch_enabled else 'off'}")
cur.execute(f"SET max_parallel_workers_per_gather = {num_workers}")
cur.execute("SET enable_bitmapscan = off")
cur.execute("SET enable_seqscan = off")
cur.execute("SET enable_indexonlyscan = off")
cur.execute("SET enable_sort = off")
apply_eviction(cur, evict_mode)
cur.execute(f'''
EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON)
SELECT length(payload) FROM prefetch_test_data
ORDER BY {column_name}
''')
result = cur.fetchone()[0]
plan = result[0]['Plan']
exec_time = result[0]['Execution Time']
blks_hit = plan.get('Shared Hit Blocks', 0)
blks_read = plan.get('Shared Read Blocks', 0)
rows = plan.get('Actual Rows', 0)
cur.execute("SHOW io_method")
io_method = cur.fetchone()[0]
cur.execute('''
INSERT INTO prefetch_test_results
(io_method, num_workers, prefetch_enabled, evict_mode, column_name,
iteration, execution_time_ms, rows_returned, blks_hit, blks_read)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
''', (io_method, num_workers, prefetch_enabled, evict_mode,
column_name, iteration, exec_time, rows, blks_hit, blks_read))
pf = 'Y' if prefetch_enabled else 'N'
print(f" [{column_name}] pf={pf} w={num_workers} evict={evict_mode}: {exec_time:.1f}ms")
def print_summary(cur):
"""Print summary of results."""
df = pd.read_sql('''
WITH avgs AS (
SELECT column_name, io_method, num_workers,
prefetch_enabled AS prefetch, evict_mode AS evict,
count(*) AS n, avg(execution_time_ms) AS ms
FROM prefetch_test_results
GROUP BY column_name, io_method, num_workers, prefetch_enabled, evict_mode
)
SELECT a.column_name, a.io_method, a.num_workers, a.evict, a.n,
round(a.ms::numeric, 1) AS off_ms,
round(b.ms::numeric, 1) AS on_ms,
round(((b.ms - a.ms) / NULLIF(b.ms + a.ms, 0) * 100)::numeric, 1) AS effect_pct
FROM avgs a JOIN avgs b USING (column_name, io_method, num_workers, evict)
WHERE NOT a.prefetch AND b.prefetch
ORDER BY column_name, io_method, num_workers, evict
''', cur.connection)
print("\n" + "=" * 80)
print("SUMMARY: Prefetch Effect (positive = slower)")
print("=" * 80)
print(df.to_string(index=False))
def main():
args = parse_args()
# Parse list arguments
columns = [c.strip() for c in args.columns.split(',')]
workers = [int(w.strip()) for w in args.workers.split(',')]
evict_modes = ['off', 'pg', 'os'] if args.evict == 'all' else [e.strip() for e in args.evict.split(',')]
print(f"Config: iterations={args.iterations}, columns={columns}, workers={workers}, evict={evict_modes}")
# Connect
connstr = f"dbname={args.dbname} host={args.host}"
if args.port:
connstr += f" port={args.port}"
conn = psycopg.connect(connstr)
conn.autocommit = True
cur = conn.cursor()
# Setup
setup_tables(cur, args.rows, args.reset)
# Show io_method
cur.execute("SHOW io_method")
io_method = cur.fetchone()[0]
print(f"\nio_method = {io_method}\n")
# Run tests
prefetch_opts = [False, True]
for i in range(1, args.iterations + 1):
print(f"Iteration {i}/{args.iterations}")
configs = list(itertools.product(columns, prefetch_opts, workers, evict_modes))
random.shuffle(configs)
for col, pf, w, evict in configs:
run_test(cur, col, pf, w, evict, i)
print()
# Summary
print_summary(cur)
conn.close()
if __name__ == "__main__":
main()