1"""Chapter 19: Search and inverted indexes. Tested with Python 3, standard library only.
2
3Running example: searching the bookstore's catalogue.
4"""
5import math
6import re
7import heapq
8from collections import Counter, defaultdict
9
10STOPWORDS = {"a", "an", "the", "of", "and", "or", "to", "in", "for", "on", "with", "is"}
11
12
13# ---------- 19.2 tokenizing and normalizing ----------
14
15def stem(word):
16 """A deliberately tiny stemmer: enough to show the idea, not for production."""
17 for suffix in ("ing", "s"):
18 if word.endswith(suffix) and len(word) - len(suffix) >= 3:
19 return word[: -len(suffix)]
20 return word
21
22
23def analyze(text):
24 """lowercase -> split on non-letters -> drop stopwords -> stem."""
25 words = re.findall(r"[a-z0-9]+", text.lower())
26 return [stem(w) for w in words if w not in STOPWORDS]
27
28
29# ---------- 19.3 inverted index ----------
30
31class InvertedIndex:
32 def __init__(self):
33 self.postings = defaultdict(dict) # term -> {doc_id: [positions]}
34 self.doc_len = {}
35 self.docs = {}
36 self.version = {} # doc_id -> last applied version (19.6)
37
38 def add(self, doc_id, text):
39 self.remove(doc_id)
40 terms = analyze(text)
41 self.docs[doc_id] = text
42 self.doc_len[doc_id] = len(terms)
43 for pos, term in enumerate(terms):
44 self.postings[term].setdefault(doc_id, []).append(pos)
45
46 def remove(self, doc_id):
47 if doc_id not in self.docs:
48 return
49 for term in set(analyze(self.docs[doc_id])):
50 self.postings[term].pop(doc_id, None)
51 if not self.postings[term]:
52 del self.postings[term]
53 del self.docs[doc_id], self.doc_len[doc_id]
54
55 def search_and(self, query):
56 terms = analyze(query)
57 if not terms:
58 return []
59 lists = sorted((set(self.postings.get(t, {})) for t in terms), key=len)
60 result = lists[0]
61 for other in lists[1:]: # start with the rarest term: smallest sets first
62 result = result & other
63 return sorted(result)
64
65 def search_phrase(self, phrase):
66 terms = analyze(phrase)
67 hits = []
68 for doc in self.search_and(phrase):
69 starts = self.postings[terms[0]][doc]
70 if any(all(p + i in self.postings[t][doc] for i, t in enumerate(terms)) for p in starts):
71 hits.append(doc)
72 return hits
73
74 # ---------- 19.4 ranking ----------
75
76 def idf(self, term):
77 n, df = len(self.docs), len(self.postings.get(term, {}))
78 return math.log(1 + (n - df + 0.5) / (df + 0.5)) # the BM25 idf form
79
80 def tfidf(self, query):
81 n = len(self.docs)
82 scores = Counter()
83 for t in analyze(query):
84 df = len(self.postings.get(t, {}))
85 if df == 0:
86 continue
87 for doc, positions in self.postings[t].items():
88 scores[doc] += len(positions) * math.log(n / df)
89 return scores
90
91 def bm25(self, query, k1=1.2, b=0.75):
92 avg = sum(self.doc_len.values()) / len(self.doc_len)
93 scores = Counter()
94 for t in analyze(query):
95 for doc, positions in self.postings.get(t, {}).items():
96 tf = len(positions)
97 norm = tf * (k1 + 1) / (tf + k1 * (1 - b + b * self.doc_len[doc] / avg))
98 scores[doc] += self.idf(t) * norm
99 return scores
100
101 def top(self, query, k=3, scorer="bm25"):
102 scores = self.bm25(query) if scorer == "bm25" else self.tfidf(query)
103 return [d for d, _ in heapq.nlargest(k, scores.items(), key=lambda x: (x[1], -x[0]))]
104
105 # ---------- 19.6 keeping in sync ----------
106
107 def apply_change(self, event):
108 """event = (doc_id, version, text or None for delete). Ignores stale events."""
109 doc_id, version, text = event
110 if version <= self.version.get(doc_id, 0):
111 return False
112 self.version[doc_id] = version
113 if text is None:
114 self.remove(doc_id)
115 else:
116 self.add(doc_id, text)
117 return True
118
119
120# ---------- 19.5 prefix search ----------
121
122class Trie:
123 def __init__(self):
124 self.root = {}
125
126 def insert(self, word, popularity):
127 node = self.root
128 for ch in word:
129 node = node.setdefault(ch, {})
130 node["$"] = popularity
131
132 def complete(self, prefix, k=3):
133 node = self.root
134 for ch in prefix:
135 if ch not in node:
136 return []
137 node = node[ch]
138 found, stack = [], [(node, prefix)]
139 while stack:
140 n, word = stack.pop()
141 for ch, child in n.items():
142 if ch == "$":
143 found.append((child, word))
144 else:
145 stack.append((child, word + ch))
146 return [w for _, w in sorted(found, key=lambda x: (-x[0], x[1]))[:k]]
147
148
149# ---------- 19.8 sharding ----------
150
151def sharded_top(shards, query, k):
152 """Document-partitioned: ask every shard for its top k, merge."""
153 merged = []
154 for shard in shards:
155 scores = shard.bm25(query) # note: each shard uses its own idf statistics
156 merged.extend(heapq.nlargest(k, scores.items(), key=lambda x: x[1]))
157 return [d for d, _ in heapq.nlargest(k, merged, key=lambda x: x[1])]
158
159
160# ---------- Level 4 ----------
161
162class BuggyIndex:
163 """Lower-cases at index time, but not at query time."""
164
165 def __init__(self):
166 self.postings = defaultdict(set)
167
168 def add(self, doc_id, text):
169 for w in analyze(text):
170 self.postings[w].add(doc_id)
171
172 def search(self, query):
173 return sorted(self.postings.get(query, set())) # bug: raw query
174
175
176def fixed_search(index, query):
177 terms = analyze(query) # same analyzer as indexing
178 if not terms:
179 return []
180 return sorted(set.intersection(*(index.postings.get(t, set()) for t in terms)))
181
182
183# ---------- checks ----------
184
185CATALOGUE = {
186 1: "Designing Data-Intensive Applications",
187 2: "Database Internals: a deep dive into storage engines",
188 3: "The Art of Computer Programming",
189 4: "Cooking with Spices: recipes from Kerala",
190 5: "System Design Interview: an insider's guide",
191 6: "Databases for beginners: databases, databases everywhere",
192 7: "Designing Distributed Systems",
193}
194
195
196def run_checks():
197 assert analyze("The Databases of Kerala!") == ["database", "kerala"]
198 assert analyze("Designing") == analyze("design")
199
200 idx = InvertedIndex()
201 for d, t in CATALOGUE.items():
202 idx.add(d, t)
203 assert idx.search_and("database") == [2, 6]
204 assert idx.search_and("designing systems") == [5, 7]
205 assert idx.search_phrase("system design") == [5]
206 assert idx.search_and("quantum") == []
207
208 # ranking: book 6 repeats "databases" 3 times in a short title; BM25 caps repetition
209 tfidf = idx.tfidf("databases")
210 bm = idx.bm25("databases")
211 assert tfidf[6] / tfidf[2] == 3.0 # TF-IDF: three times the score
212 assert 1.0 < bm[6] / bm[2] < 2.0 # BM25: under twice
213 assert idx.top("database internals") == [2, 6]
214
215 # trie autocomplete by popularity
216 trie = Trie()
217 for word, pop in [("database", 900), ("data", 1200), ("datastore", 50),
218 ("design", 800), ("date", 300)]:
219 trie.insert(word, pop)
220 assert trie.complete("dat") == ["data", "database", "date"]
221 assert trie.complete("des") == ["design"]
222 assert trie.complete("x") == []
223
224 # sync: out-of-order change events must not resurrect old data
225 sync = InvertedIndex()
226 events = [(10, 1, "Old title about cooking"),
227 (10, 3, "New title about databases"),
228 (10, 2, "Middle title about gardening")] # arrives late
229 applied = [sync.apply_change(e) for e in events]
230 assert applied == [True, True, False]
231 assert sync.search_and("databases") == [10] and sync.search_and("gardening") == []
232 sync.apply_change((10, 4, None)) # delete
233 assert sync.search_and("databases") == []
234
235 # sharding: split the catalogue by doc id, merge top-k
236 shards = [InvertedIndex(), InvertedIndex()]
237 for d, t in CATALOGUE.items():
238 shards[d % 2].add(d, t)
239 assert set(sharded_top(shards, "designing", 2)) == {1, 7}
240
241 # Level 4
242 buggy = BuggyIndex()
243 for d, t in CATALOGUE.items():
244 buggy.add(d, t)
245 assert buggy.search("Database") == [] # the complaint
246 assert fixed_search(buggy, "Database") == [2, 6]
247
248 # Level 2 arithmetic
249 rows, row_bytes, disk_mb_s = 2_000_000, 1_000, 500
250 scan_seconds = rows * row_bytes / (disk_mb_s * 1_000_000)
251 assert scan_seconds == 4.0
252 assert 2_000_000 * 200 == 400_000_000 # word occurrences to index
253
254 return {"tfidf_ratio": tfidf[6] / tfidf[2], "bm25_ratio": round(bm[6] / bm[2], 2),
255 "complete": trie.complete("dat"), "scan_s": scan_seconds}
256
257
258if __name__ == "__main__":
259 print("All checks pass.", run_checks())