1"""Chapter 9 practice: what an index changes, shown with sqlite3's real query planner
2(Level 3 Build), and the 'function on an indexed column' bug (Level 4 Debug).
3Standard library only. Run: python3 ch09_indexes.py"""
4import bisect
5import random
6import sqlite3
7import time
8
9N_ORDERS = 200_000
10
11
12def make_db(seed=1):
13 rng = random.Random(seed)
14 db = sqlite3.connect(":memory:")
15 db.execute("""CREATE TABLE orders (
16 id INTEGER PRIMARY KEY, user_id INTEGER, created_at INTEGER,
17 status TEXT, total INTEGER)""")
18 db.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, email TEXT)")
19 db.executemany("INSERT INTO orders VALUES (?, ?, ?, ?, ?)", (
20 (i, rng.randrange(20_000), 1_700_000_000 + i * 60,
21 rng.choice(["paid", "shipped", "delivered"]), rng.randrange(100, 5000))
22 for i in range(N_ORDERS)))
23 db.executemany("INSERT INTO users VALUES (?, ?)",
24 ((u, f"User{u}@Example.com") for u in range(20_000)))
25 db.commit()
26 return db
27
28
29def plan(db, sql, params=()):
30 """The planner's chosen strategy, e.g. 'SCAN orders' or 'SEARCH orders USING INDEX ...'."""
31 return " | ".join(row[-1] for row in db.execute("EXPLAIN QUERY PLAN " + sql, params))
32
33
34def timed(db, sql, params=(), repeat=20):
35 start = time.perf_counter()
36 for _ in range(repeat):
37 db.execute(sql, params).fetchall()
38 return (time.perf_counter() - start) / repeat * 1000 # ms per query
39
40
41# ---- A tiny "index": a sorted list you can binary-search --------------------------------
42
43def scan_lookup(rows, key):
44 comparisons = 0
45 for k, rowid in rows:
46 comparisons += 1
47 if k == key:
48 return rowid, comparisons
49 return None, comparisons
50
51
52def index_lookup(sorted_keys, sorted_rowids, key):
53 """Binary search: about log2(n) comparisons instead of n."""
54 lo, hi, comparisons = 0, len(sorted_keys), 0
55 while lo < hi:
56 comparisons += 1
57 mid = (lo + hi) // 2
58 if sorted_keys[mid] < key:
59 lo = mid + 1
60 else:
61 hi = mid
62 found = lo < len(sorted_keys) and sorted_keys[lo] == key
63 return (sorted_rowids[lo] if found else None), comparisons
64
65
66def insert_time(n_indexes, rows=20_000):
67 db = sqlite3.connect(":memory:")
68 db.execute("CREATE TABLE t (a INTEGER, b INTEGER, c INTEGER, d INTEGER, e INTEGER)")
69 for i, col in enumerate("abcde"[:n_indexes]):
70 db.execute(f"CREATE INDEX idx_{col} ON t({col})")
71 rng = random.Random(2)
72 data = [tuple(rng.randrange(1_000_000) for _ in range(5)) for _ in range(rows)]
73 start = time.perf_counter()
74 db.executemany("INSERT INTO t VALUES (?, ?, ?, ?, ?)", data)
75 db.commit()
76 return time.perf_counter() - start
77
78
79def run_checks():
80 db = make_db()
81 q_user = "SELECT id, total FROM orders WHERE user_id = ?"
82
83 # 1. Without an index: the planner scans every row.
84 assert plan(db, q_user, (7,)).startswith("SCAN orders")
85 slow = timed(db, q_user, (7,))
86
87 # 2. With an index: it searches.
88 db.execute("CREATE INDEX idx_orders_user ON orders(user_id)")
89 assert "SEARCH orders USING INDEX idx_orders_user" in plan(db, q_user, (7,))
90 fast = timed(db, q_user, (7,))
91 assert fast * 5 < slow, (slow, fast)
92
93 # 3. Composite index and the leftmost-prefix rule.
94 db.execute("DROP INDEX idx_orders_user")
95 db.execute("CREATE INDEX idx_user_time ON orders(user_id, created_at)")
96 assert "USING INDEX idx_user_time (user_id=?)" in plan(db, q_user, (7,))
97 assert "(user_id=? AND created_at>?)" in plan(
98 db, "SELECT id FROM orders WHERE user_id = ? AND created_at > ?", (7, 0))
99 only_time = plan(db, "SELECT id FROM orders WHERE created_at > ?", (1_700_000_000,))
100 assert "idx_user_time" not in only_time or "SCAN" in only_time # can't seek on 2nd column alone
101
102 # 4. Covering index: everything the query needs is inside the index.
103 cov = plan(db, "SELECT created_at FROM orders WHERE user_id = ?", (7,))
104 assert "COVERING INDEX" in cov, cov
105
106 # 5. Index on a low-selectivity column: the planner may still choose it, but it reads ~1/3 of rows.
107 # (Shown in the chapter text; not asserted because planners differ.)
108
109 # 6. The toy index: binary search vs scan.
110 rng = random.Random(5)
111 keys = rng.sample(range(10_000_000), 1_000_000)
112 rows = list(zip(keys, range(len(keys))))
113 pairs = sorted(rows)
114 sk, sr = [k for k, _ in pairs], [r for _, r in pairs]
115 target = keys[-1]
116 _, scan_cmp = scan_lookup(rows, target)
117 _, idx_cmp = index_lookup(sk, sr, target)
118 assert scan_cmp == 1_000_000 and idx_cmp <= 21
119
120 # 7. Write cost: more indexes, slower inserts.
121 t0, t5 = insert_time(0), insert_time(5)
122 assert t5 > t0
123
124 # Level 4: a function on the indexed column hides the index.
125 db.execute("CREATE INDEX idx_users_email ON users(email)")
126 bad = plan(db, "SELECT id FROM users WHERE lower(email) = ?", ("user7@example.com",))
127 assert bad.startswith("SCAN users"), bad
128 db.execute("CREATE INDEX idx_users_email_lower ON users(lower(email))") # expression index
129 good = plan(db, "SELECT id FROM users WHERE lower(email) = ?", ("user7@example.com",))
130 assert "SEARCH users USING INDEX idx_users_email_lower" in good, good
131
132 print(f"All checks pass. user_id lookup: {slow:.2f} ms scanning vs {fast:.3f} ms with an index.")
133 print(f"Toy index: {scan_cmp:,} comparisons scanning vs {idx_cmp} with binary search.")
134 print(f"20,000 inserts: {t0 * 1000:.0f} ms with no indexes vs {t5 * 1000:.0f} ms with 5.")
135 print("Plans:", "\n ", bad, "\n ", good, "\n ", cov)
136
137
138if __name__ == "__main__":
139 run_checks()