lib/sql/src/search/rank.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const sql = @import("../root.zig");
3 const engine_mod = @import("engine.zig");
4 const hit_mod = @import("hit.zig");
5 const posting_mod = @import("posting.zig");
6 const query_mod = @import("query.zig");
7 const token_mod = @import("token.zig");
8
9 const Allocator = std.mem.Allocator;
10
11 pub const bm25_min_idf = 0.000001;
12
13 pub const RankingOptions = struct {
14 k1: f64 = 1.2,
15 b: f64 = 0.75,
16 };
17
18 pub const DocumentLengthCursor = struct {
19 allocator: Allocator,
20 segments: []posting_mod.LengthSegmentCursor,
21 direct: []posting_mod.DirectDocumentLength,
22 direct_index: usize = 0,
23 rowid: ?i64 = null,
24 length: usize = 0,
25
26 pub fn init(search: anytype, allocator: Allocator) engine_mod.Error!DocumentLengthCursor {
27 var lookup: sql.IndexScan = undefined;
28 try search.terms.lookupPayloads(&lookup, allocator, &posting_mod.documentLengthValues());
29 defer lookup.deinit();
30 var segments: std.ArrayList(posting_mod.LengthSegmentCursor) = .empty;
31 var direct: std.ArrayList(posting_mod.DirectDocumentLength) = .empty;
32 errdefer {
33 for (segments.items) |*segment| segment.deinit(allocator);
34 segments.deinit(allocator);
35 direct.deinit(allocator);
36 }
37 while (try lookup.next()) |entry| {
38 if (posting_mod.lengthSegmentPayloadCount(entry.payload)) |count| {
39 try segments.append(allocator, .{
40 .payload = try allocator.dupe(u8, entry.payload),
41 .count = count,
42 });
43 } else {
44 try direct.append(allocator, .{
45 .rowid = entry.rowid,
46 .length = try posting_mod.documentLengthPayloadMaybe(entry.payload),
47 });
48 }
49 }
50 std.mem.sort(posting_mod.DirectDocumentLength, direct.items, {}, posting_mod.directDocumentLengthLess);
51 var cursor = DocumentLengthCursor{
52 .allocator = allocator,
53 .segments = try segments.toOwnedSlice(allocator),
54 .direct = try direct.toOwnedSlice(allocator),
55 };
56 try cursor.advance();
57 return cursor;
58 }
59
60 pub fn deinit(self: *DocumentLengthCursor) void {
61 for (self.segments) |*segment| segment.deinit(self.allocator);
62 self.allocator.free(self.segments);
63 self.allocator.free(self.direct);
64 self.* = undefined;
65 }
66
67 pub fn advance(self: *DocumentLengthCursor) engine_mod.Error!void {
68 while (true) {
69 const rowid = try self.nextRowid() orelse {
70 self.rowid = null;
71 self.length = 0;
72 return;
73 };
74 var segment_length: ?usize = null;
75 for (self.segments) |*segment| {
76 while (true) {
77 const length = (try segment.current()) orelse break;
78 if (length.rowid != rowid) break;
79 if (segment_length != null) return error.InvalidSearchIndex;
80 segment_length = length.length;
81 segment.advance();
82 }
83 }
84 var direct_length: ??usize = null;
85 while (self.direct_index < self.direct.len and self.direct[self.direct_index].rowid == rowid) : (self.direct_index += 1) {
86 direct_length = self.direct[self.direct_index].length;
87 }
88 const maybe_length = if (direct_length) |length| length else segment_length;
89 const length = maybe_length orelse continue;
90 self.rowid = rowid;
91 self.length = length;
92 return;
93 }
94 }
95
96 pub fn lengthFor(self: *DocumentLengthCursor, rowid: i64) engine_mod.Error!usize {
97 while (self.rowid) |current| {
98 if (current == rowid) return self.length;
99 if (current > rowid) return error.InvalidSearchIndex;
100 try self.advance();
101 }
102 return error.InvalidSearchIndex;
103 }
104
105 pub fn nextRowid(self: *const DocumentLengthCursor) engine_mod.Error!?i64 {
106 var rowid: ?i64 = null;
107 if (self.direct_index < self.direct.len) rowid = self.direct[self.direct_index].rowid;
108 for (self.segments) |segment| {
109 const length = (try segment.current()) orelse continue;
110 if (rowid == null or length.rowid < rowid.?) rowid = length.rowid;
111 }
112 return rowid;
113 }
114 };
115
116 pub const RankingTerm = struct {
117 atom: *const query_mod.Atom,
118 posting_hits: []hit_mod.Hit,
119 document_frequency: usize,
120 idf: f64,
121
122 pub fn deinit(self: *RankingTerm, allocator: Allocator) void {
123 if (self.posting_hits.len != 0) allocator.free(self.posting_hits);
124 self.* = undefined;
125 }
126 };
127
128 pub const SingleAtomCandidate = struct {
129 rowid: i64,
130 length: usize,
131 frequency: usize,
132 };
133
134 pub const ExactRankingCursor = struct {
135 cursor: ExactTermCursor,
136 idf: f64,
137
138 pub fn deinit(self: *ExactRankingCursor) void {
139 self.cursor.deinit();
140 self.* = undefined;
141 }
142 };
143
144 pub const ExactTermCursor = struct {
145 allocator: Allocator,
146 segments: []posting_mod.SegmentCursor,
147 direct: []posting_mod.DirectPosting,
148 direct_index: usize = 0,
149 current: ?hit_mod.Hit = null,
150
151 pub fn init(search: anytype, allocator: Allocator, term: []const u8) engine_mod.Error!ExactTermCursor {
152 var key_buffer: [token_mod.max_token_bytes + 1]u8 = undefined;
153 const key_text = posting_mod.postingKeyBuffer(&key_buffer, '=', term);
154 return try ExactTermCursor.initKey(search, allocator, key_text);
155 }
156
157 pub fn initPhrasePair(search: anytype, allocator: Allocator, first: []const u8, second: []const u8) engine_mod.Error!ExactTermCursor {
158 var key_buffer: [posting_mod.phrase_pair_posting_key_max_bytes]u8 = undefined;
159 const key_text = posting_mod.phrasePairPostingKeyBuffer(&key_buffer, first, second);
160 return try ExactTermCursor.initKey(search, allocator, key_text);
161 }
162
163 pub fn initKey(search: anytype, allocator: Allocator, key_text: []const u8) engine_mod.Error!ExactTermCursor {
164 var lookup: sql.IndexScan = undefined;
165 try search.terms.lookupPayloads(&lookup, allocator, &.{.{ .text = key_text }});
166 defer lookup.deinit();
167 var segments: std.ArrayList(posting_mod.SegmentCursor) = .empty;
168 var direct: std.ArrayList(posting_mod.DirectPosting) = .empty;
169 errdefer {
170 for (segments.items) |*segment| segment.deinit(allocator);
171 segments.deinit(allocator);
172 direct.deinit(allocator);
173 }
174 while (try lookup.next()) |entry| {
175 if (posting_mod.segmentPayloadCount(entry.payload)) |count| {
176 try segments.ensureUnusedCapacity(allocator, 1);
177 segments.appendAssumeCapacity(try posting_mod.SegmentCursor.init(allocator, entry.payload, count));
178 } else {
179 try direct.append(allocator, .{
180 .rowid = entry.rowid,
181 .score = try posting_mod.postingPayloadCount(entry.payload),
182 });
183 }
184 }
185 std.mem.sort(posting_mod.DirectPosting, direct.items, {}, posting_mod.directPostingLess);
186 return .{
187 .allocator = allocator,
188 .segments = try segments.toOwnedSlice(allocator),
189 .direct = try direct.toOwnedSlice(allocator),
190 };
191 }
192
193 pub fn deinit(self: *ExactTermCursor) void {
194 for (self.segments) |*segment| segment.deinit(self.allocator);
195 self.allocator.free(self.segments);
196 self.allocator.free(self.direct);
197 self.* = undefined;
198 }
199
200 pub fn advance(self: *ExactTermCursor) engine_mod.Error!void {
201 while (true) {
202 const rowid = try self.nextRowid() orelse {
203 self.current = null;
204 return;
205 };
206 var segment_score: usize = 0;
207 for (self.segments) |*segment| {
208 while (true) {
209 const hit = (try segment.current()) orelse break;
210 if (hit.rowid != rowid) break;
211 segment_score += hit.score;
212 try segment.advance();
213 }
214 }
215 var direct_score: ?usize = null;
216 while (self.direct_index < self.direct.len and self.direct[self.direct_index].rowid == rowid) : (self.direct_index += 1) {
217 direct_score = self.direct[self.direct_index].score;
218 }
219 const score = direct_score orelse segment_score;
220 if (score == 0) continue;
221 self.current = .{ .rowid = rowid, .score = score };
222 return;
223 }
224 }
225
226 pub fn documentFrequency(self: *ExactTermCursor) engine_mod.Error!usize {
227 if (try self.segmentDocumentFrequency()) |count| return count;
228 self.reset();
229 var count: usize = 0;
230 while (true) {
231 try self.advance();
232 if (self.current == null) break;
233 count += 1;
234 }
235 self.reset();
236 return count;
237 }
238
239 pub fn segmentDocumentFrequency(self: *const ExactTermCursor) engine_mod.Error!?usize {
240 if (self.direct.len != 0) return null;
241 var count: usize = 0;
242 var previous_last: ?i64 = null;
243 for (self.segments) |segment| {
244 if (segment.count == 0) return error.InvalidSearchIndex;
245 const bounds = try posting_mod.segmentPayloadBounds(segment.payload);
246 if (previous_last) |rowid| {
247 if (bounds.first <= rowid) return null;
248 }
249 count = std.math.add(usize, count, segment.count) catch return error.InvalidSearchIndex;
250 previous_last = bounds.last;
251 }
252 return count;
253 }
254
255 pub fn reset(self: *ExactTermCursor) void {
256 for (self.segments) |*segment| segment.reset();
257 self.direct_index = 0;
258 self.current = null;
259 }
260
261 pub fn nextRowid(self: *ExactTermCursor) engine_mod.Error!?i64 {
262 var rowid: ?i64 = null;
263 if (self.direct_index < self.direct.len) rowid = self.direct[self.direct_index].rowid;
264 for (self.segments) |*segment| {
265 const hit = (try segment.current()) orelse continue;
266 if (rowid == null or hit.rowid < rowid.?) rowid = hit.rowid;
267 }
268 return rowid;
269 }
270 };
271
272 pub fn appendSingleAtomCandidate(
273 allocator: Allocator,
274 candidates: *std.ArrayList(SingleAtomCandidate),
275 lengths: *DocumentLengthCursor,
276 document_frequency: *usize,
277 rowid: i64,
278 frequency: usize,
279 ) engine_mod.Error!void {
280 if (frequency == 0) return;
281 document_frequency.* += 1;
282 try candidates.append(allocator, .{
283 .rowid = rowid,
284 .length = try lengths.lengthFor(rowid),
285 .frequency = frequency,
286 });
287 }
288
289 pub fn appendPrefixRankedHit(
290 allocator: Allocator,
291 ranked: *std.ArrayList(hit_mod.RankedHit),
292 lengths: *DocumentLengthCursor,
293 rowid: i64,
294 frequency: usize,
295 idf: f64,
296 stats: posting_mod.CorpusStats,
297 options: RankingOptions,
298 limit: usize,
299 ) engine_mod.Error!void {
300 if (frequency == 0) return;
301 try hit_mod.appendRankedHit(allocator, ranked, limit, .{
302 .rowid = rowid,
303 .rank = -bm25Score(idf, frequency, try lengths.lengthFor(rowid), stats, options),
304 });
305 }
306
307 pub fn nextExactOrHit(cursors: []ExactTermCursor) engine_mod.Error!?hit_mod.Hit {
308 var min_rowid: ?i64 = null;
309 for (cursors) |cursor| {
310 const hit = cursor.current orelse continue;
311 if (min_rowid == null or hit.rowid < min_rowid.?) min_rowid = hit.rowid;
312 }
313 const rowid = min_rowid orelse return null;
314 var score: usize = 0;
315 for (cursors) |*cursor| {
316 const hit = cursor.current orelse continue;
317 if (hit.rowid != rowid) continue;
318 score += hit.score;
319 try cursor.advance();
320 }
321 return .{ .rowid = rowid, .score = score };
322 }
323
324 pub fn nextPhrasePairCandidateRow(left: *ExactTermCursor, right: *ExactTermCursor) engine_mod.Error!?i64 {
325 while (left.current != null and right.current != null) {
326 const left_rowid = left.current.?.rowid;
327 const right_rowid = right.current.?.rowid;
328 if (left_rowid == right_rowid) {
329 try left.advance();
330 try right.advance();
331 return left_rowid;
332 }
333 if (left_rowid < right_rowid) {
334 try left.advance();
335 } else {
336 try right.advance();
337 }
338 }
339 return null;
340 }
341
342 pub fn postingCursorFrequencyAt(cursor: *ExactTermCursor, rowid: i64) engine_mod.Error!usize {
343 while (cursor.current) |hit| {
344 if (hit.rowid < rowid) {
345 try cursor.advance();
346 continue;
347 }
348 if (hit.rowid > rowid) return 0;
349 const score = hit.score;
350 try cursor.advance();
351 return score;
352 }
353 return 0;
354 }
355
356 pub fn nextExactRankedOrHit(cursors: []ExactRankingCursor, lengths: *DocumentLengthCursor, stats: posting_mod.CorpusStats, options: RankingOptions) engine_mod.Error!?hit_mod.RankedHit {
357 var min_rowid: ?i64 = null;
358 for (cursors) |cursor| {
359 const hit = cursor.cursor.current orelse continue;
360 if (min_rowid == null or hit.rowid < min_rowid.?) min_rowid = hit.rowid;
361 }
362 const rowid = min_rowid orelse return null;
363 const length = try lengths.lengthFor(rowid);
364 var score: f64 = 0.0;
365 for (cursors) |*cursor| {
366 const hit = cursor.cursor.current orelse continue;
367 if (hit.rowid != rowid) continue;
368 score += bm25Score(cursor.idf, hit.score, length, stats, options);
369 try cursor.cursor.advance();
370 }
371 return .{ .rowid = rowid, .rank = -score };
372 }
373
374 pub fn bm25Idf(document_count: usize, document_frequency: usize) engine_mod.Error!f64 {
375 if (document_frequency > document_count) return error.InvalidSearchIndex;
376 const documents: f64 = @floatFromInt(document_count);
377 const frequency: f64 = @floatFromInt(document_frequency);
378 const value = @log((documents - frequency + 0.5) / (frequency + 0.5));
379 if (value <= bm25_min_idf) return bm25_min_idf;
380 return value;
381 }
382
383 pub fn bm25Score(idf: f64, frequency: usize, length: usize, stats: posting_mod.CorpusStats, options: RankingOptions) f64 {
384 if (frequency == 0 or stats.documents == 0 or stats.total_tokens == 0) return 0.0;
385 const term_frequency: f64 = @floatFromInt(frequency);
386 const document_length: f64 = @floatFromInt(length);
387 const average_length = @as(f64, @floatFromInt(stats.total_tokens)) / @as(f64, @floatFromInt(stats.documents));
388 const denominator = term_frequency + options.k1 * (1.0 - options.b + options.b * document_length / average_length);
389 if (denominator == 0.0) return 0.0;
390 return idf * (term_frequency * (options.k1 + 1.0)) / denominator;
391 }
392
393 pub fn documentLengthFromSorted(lengths: []const posting_mod.DocumentLength, rowid: i64) ?usize {
394 var start: usize = 0;
395 var end = lengths.len;
396 while (start < end) {
397 const middle = start + (end - start) / 2;
398 if (lengths[middle].rowid == rowid) return lengths[middle].length;
399 if (lengths[middle].rowid < rowid) {
400 start = middle + 1;
401 } else {
402 end = middle;
403 }
404 }
405 return null;
406 }
407
408 pub fn deinitRankingTerms(allocator: Allocator, terms: []RankingTerm) void {
409 for (terms) |*term| term.deinit(allocator);
410 if (terms.len != 0) allocator.free(terms);
411 }