1"""
2Run a benchmark to compare different key strategies in SQLite.
3
4In particular, does it make sense to store text as raw strings in SQLite, or
5to hash the strings as MD5/xxHash before insertion?
6"""
7
8from collections.abc import Callable
9from collections import defaultdict, namedtuple
10from contextlib import closing
11import csv
12from functools import cache
13import hashlib
14from pathlib import Path
15import random
16import sqlite3
17import statistics
18import string
19import time
20
21import xxhash
22
23
24def random_string(length: int) -> str:
25    """
26    Generate a random alphanumeric string of a specific length.
27    """
28    return "".join(random.choices(string.ascii_letters + string.digits, k=length))
29
30
31type KeyStrategy = Callable[[str], str]
32
33
34def identity(key: str) -> str:
35    """Return a key unmodified."""
36    return key
37
38
39def uncached_md5(s: str) -> str:
40    """Return the MD5 hex digest of a string, computed fresh each time."""
41    h = hashlib.md5()
42    h.update(s.encode("utf8"))
43    return h.hexdigest()
44
45
46@cache
47def cached_md5(s: str) -> str:
48    """Return the MD5 hex digest of a string, cached after first hash."""
49    return uncached_md5(s)
50
51
52def uncached_xxh64(s: str) -> str:
53    """Return the XXH64 hex digest of a string, computed fresh each time."""
54    return xxhash.xxh64_hexdigest(s.encode("utf8"))
55
56
57@cache
58def cached_xxh64(s: str) -> str:
59    """Return the XXH64 hex digest of a string, cached after first hash."""
60    return uncached_xxh64(s)
61
62
63Result = namedtuple("Result", ["insert_duration", "lookup_duration"])
64
65
66def run_benchmark(
67    key_size: int, key_strategy: KeyStrategy, *, num_operations: int
68) -> Result:
69    """
70    Run the benchmark for a single key size.
71    """
72    # Use an in-memory SQLite database to isolate CPU/B-tree performance
73    # from disk I/O.
74    #
75    # Disk performance would likely also have an effect if a key spans
76    # multiple pages, but it would only exaggerate the result.
77    with closing(sqlite3.connect(":memory:")) as con:
78        # Add the same settings and table as I'm using in SQLiteStorage.
79        con.execute("PRAGMA journal_mode = WAL;")
80        con.execute("PRAGMA synchronous = NORMAL;")
81        con.execute("CREATE TABLE cache(key TEXT PRIMARY KEY, value TEXT)")
82        con.commit()
83
84        # Generate a batch of unique keys and values upfront
85        keys = [random_string(key_size) for _ in range(num_operations)]
86        values = [random_string(10) for _ in range(num_operations)]
87
88        expected = {k: v for k, v in zip(keys, values)}
89
90        # 1. Benchmark insertions
91        start_time = time.perf_counter()
92        cur = con.cursor()
93        for k, v in zip(keys, values):
94            insert_key = key_strategy(k)
95            cur.execute("INSERT INTO cache VALUES (?, ?)", (insert_key, v))
96        con.commit()
97        insert_duration = time.perf_counter() - start_time
98
99        # 2. Benchmark lookups
100        # Shuffle keys to simulate random cache access patterns
101        shuffled_keys = list(keys)
102        random.shuffle(shuffled_keys)
103
104        start_time = time.perf_counter()
105        cur = con.cursor()
106        for k in shuffled_keys:
107            lookup_key = key_strategy(k)
108            cur.execute("SELECT value FROM cache WHERE key = ?", (lookup_key,))
109            assert cur.fetchone() == (expected[k],), (cur.fetchone(), expected[k])
110        lookup_duration = time.perf_counter() - start_time
111
112    return Result(insert_duration, lookup_duration)
113
114
115if __name__ == "__main__":
116    sizes = [16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
117    all_results = defaultdict(lambda: {sz: [] for sz in sizes})
118
119    key_strategies = [
120        identity,
121        uncached_md5,
122        cached_md5,
123        uncached_xxh64,
124        cached_xxh64,
125    ]
126
127    # Run six interleaved samples and average the results to reduce
128    # the noise of local CPU/memory contention.
129    for i in range(1, 7):
130        print(f"Getting sample {i}...")
131        for sz in sizes:
132            for ks in key_strategies:
133                all_results[ks][sz].append(run_benchmark(sz, ks, num_operations=5000))
134
135    results = defaultdict(dict)
136    for sz in sizes:
137        for ks in key_strategies:
138            results[ks][sz] = Result(
139                statistics.mean(ins for ins, _ in all_results[ks][sz]),
140                statistics.mean(lok for _, lok in all_results[ks][sz]),
141            )
142
143    from pprint import pprint
144
145    pprint(results)
146
147    results_path = Path(__file__).parent / (Path(__file__).stem + ".csv")
148
149    with open(results_path, "w") as f:
150        fieldnames = ["size"]
151        for ks in key_strategies:
152            fieldnames.append(f"{ks.__name__} (insert)")
153            fieldnames.append(f"{ks.__name__} (lookup)")
154
155        writer = csv.DictWriter(f, fieldnames=fieldnames)
156        writer.writeheader()
157
158        for sz in sizes:
159            row = {
160                "size": sz,
161            }
162            for ks in key_strategies:
163                row[f"{ks.__name__} (insert)"] = results[ks][sz].insert_duration
164                row[f"{ks.__name__} (lookup)"] = results[ks][sz].lookup_duration
165            writer.writerow(row)