cr-search/benchmark.py

import json
import shutil
from statistics import mean, stdev

from alive_progress import alive_bar
from peewee import SelectQuery
from psycopg2._psycopg import cursor

from app import db
from server import search, suggest, exact_search


def benchmark_query(query: SelectQuery, filename: str = None) -> tuple[float, float]:
    query, params = query.sql()

    query = "EXPLAIN (ANALYZE, COSTS, VERBOSE, BUFFERS, FORMAT JSON) " + query

    cur: cursor = db.execute_sql(query, params=params)

    result = cur.fetchone()[0][0]
    if filename:
        with open(f"benchmark/{filename}.json", "w") as f:
            json.dump(result, f, indent=2)
        with open(f"benchmark/{filename}.txt", "w") as f:
            f.write(query)

    return result["Planning Time"], result["Execution Time"]


def statistics(query: SelectQuery, filename: str, repeats: int = 500) -> None:
    ts = shutil.get_terminal_size((80, 20))
    print(filename.center(ts.columns, "-"))
    planning_times = []
    execution_times = []
    benchmark_query(query, filename=filename)

    with alive_bar(repeats) as bar:
        for i in range(repeats):
            plantime, exetime = benchmark_query(query)
            planning_times.append(plantime)
            execution_times.append(exetime)
            bar()
    print(mean(planning_times), stdev(planning_times))
    print(mean(execution_times), stdev(execution_times))


test_search = search("hello", 1000, "campaign2", 200)
statistics(test_search, filename="search_hello")
test_search = exact_search("hello", 1000, "campaign2", 200)
statistics(test_search, filename="exact_search", repeats=50)
test_search = search("a very long search query with a lot of stop word", 1000, "campaign2", 200)
statistics(test_search, filename="search_long")

test_search = suggest("gnoll", 1000, "campaign2")
statistics(test_search, filename="suggest_simple")
test_search = suggest("gu", 1000, "campaign2")
statistics(test_search, filename="suggest_two_letter", repeats=100)