lib/sql/src/search/hit.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const engine_mod = @import("engine.zig");
3 const posting_mod = @import("posting.zig");
4
5 const Allocator = std.mem.Allocator;
6
7 pub const Hit = struct {
8 rowid: i64,
9 score: usize,
10 };
11
12 pub const RankedHit = struct {
13 rowid: i64,
14 rank: f64,
15 };
16
17 pub const HitSet = struct {
18 hits: []Hit,
19 total: usize,
20 };
21
22 pub const RawHit = struct {
23 rowid: i64,
24 score: usize,
25 direct: bool,
26 };
27
28 pub fn appendRawHits(allocator: Allocator, hits: *std.ArrayList(RawHit), rowid: i64, payload: []const u8) engine_mod.Error!void {
29 if (posting_mod.segmentPayloadCount(payload)) |count| {
30 try hits.ensureUnusedCapacity(allocator, count);
31 var decoder = try posting_mod.SegmentDecoder.initValidated(payload, count);
32 while (try decoder.next()) |hit| {
33 hits.appendAssumeCapacity(.{
34 .rowid = hit.rowid,
35 .score = hit.score,
36 .direct = false,
37 });
38 }
39 return;
40 }
41 try hits.append(allocator, .{
42 .rowid = rowid,
43 .score = try posting_mod.postingPayloadCount(payload),
44 .direct = true,
45 });
46 }
47
48 pub fn coalescedHits(allocator: Allocator, raw: *std.ArrayList(RawHit)) engine_mod.Error![]Hit {
49 defer raw.deinit(allocator);
50 if (raw.items.len == 0) return &.{};
51 std.mem.sort(RawHit, raw.items, {}, rawHitLess);
52 var hits: std.ArrayList(Hit) = .empty;
53 errdefer hits.deinit(allocator);
54 var index: usize = 0;
55 while (index < raw.items.len) {
56 const rowid = raw.items[index].rowid;
57 var segment_score: usize = 0;
58 var direct_score: ?usize = null;
59 while (index < raw.items.len and raw.items[index].rowid == rowid) : (index += 1) {
60 const hit = raw.items[index];
61 if (hit.direct) {
62 direct_score = hit.score;
63 } else {
64 segment_score += hit.score;
65 }
66 }
67 const score = direct_score orelse segment_score;
68 if (score != 0) try hits.append(allocator, .{ .rowid = rowid, .score = score });
69 }
70 return try hits.toOwnedSlice(allocator);
71 }
72
73 pub fn rawHitLess(_: void, left: RawHit, right: RawHit) bool {
74 if (left.rowid != right.rowid) return left.rowid < right.rowid;
75 if (left.direct != right.direct) return !left.direct and right.direct;
76 return left.score < right.score;
77 }
78
79 pub fn intersectRows(allocator: Allocator, left: []const i64, right: []const i64) engine_mod.Error![]i64 {
80 var out: std.ArrayList(i64) = .empty;
81 errdefer out.deinit(allocator);
82 var left_index: usize = 0;
83 var right_index: usize = 0;
84 while (left_index < left.len and right_index < right.len) {
85 if (left[left_index] == right[right_index]) {
86 try out.append(allocator, left[left_index]);
87 left_index += 1;
88 right_index += 1;
89 } else if (left[left_index] < right[right_index]) {
90 left_index += 1;
91 } else {
92 right_index += 1;
93 }
94 }
95 return try out.toOwnedSlice(allocator);
96 }
97
98 pub fn subtractRows(allocator: Allocator, left: []const i64, right: []const i64) engine_mod.Error![]i64 {
99 var out: std.ArrayList(i64) = .empty;
100 errdefer out.deinit(allocator);
101 var left_index: usize = 0;
102 var right_index: usize = 0;
103 while (left_index < left.len) {
104 while (right_index < right.len and right[right_index] < left[left_index]) right_index += 1;
105 if (right_index < right.len and right[right_index] == left[left_index]) {
106 left_index += 1;
107 continue;
108 }
109 try out.append(allocator, left[left_index]);
110 left_index += 1;
111 }
112 return try out.toOwnedSlice(allocator);
113 }
114
115 pub fn unionRows(allocator: Allocator, left: []const i64, right: []const i64) engine_mod.Error![]i64 {
116 var out: std.ArrayList(i64) = .empty;
117 errdefer out.deinit(allocator);
118 try out.ensureTotalCapacity(allocator, left.len + right.len);
119 var left_index: usize = 0;
120 var right_index: usize = 0;
121 while (left_index < left.len and right_index < right.len) {
122 if (left[left_index] == right[right_index]) {
123 out.appendAssumeCapacity(left[left_index]);
124 left_index += 1;
125 right_index += 1;
126 } else if (left[left_index] < right[right_index]) {
127 out.appendAssumeCapacity(left[left_index]);
128 left_index += 1;
129 } else {
130 out.appendAssumeCapacity(right[right_index]);
131 right_index += 1;
132 }
133 }
134 while (left_index < left.len) : (left_index += 1) out.appendAssumeCapacity(left[left_index]);
135 while (right_index < right.len) : (right_index += 1) out.appendAssumeCapacity(right[right_index]);
136 return try out.toOwnedSlice(allocator);
137 }
138
139 pub fn intersectHits(allocator: Allocator, left: []const Hit, right: []const Hit) engine_mod.Error![]Hit {
140 var out: std.ArrayList(Hit) = .empty;
141 errdefer out.deinit(allocator);
142 var left_index: usize = 0;
143 var right_index: usize = 0;
144 while (left_index < left.len and right_index < right.len) {
145 if (left[left_index].rowid == right[right_index].rowid) {
146 try out.append(allocator, .{
147 .rowid = left[left_index].rowid,
148 .score = left[left_index].score + right[right_index].score,
149 });
150 left_index += 1;
151 right_index += 1;
152 } else if (left[left_index].rowid < right[right_index].rowid) {
153 left_index += 1;
154 } else {
155 right_index += 1;
156 }
157 }
158 return try out.toOwnedSlice(allocator);
159 }
160
161 pub fn subtractHits(allocator: Allocator, left: []const Hit, right: []const i64) engine_mod.Error![]Hit {
162 var out: std.ArrayList(Hit) = .empty;
163 errdefer out.deinit(allocator);
164 var left_index: usize = 0;
165 var right_index: usize = 0;
166 while (left_index < left.len) {
167 while (right_index < right.len and right[right_index] < left[left_index].rowid) right_index += 1;
168 if (right_index < right.len and right[right_index] == left[left_index].rowid) {
169 left_index += 1;
170 continue;
171 }
172 try out.append(allocator, left[left_index]);
173 left_index += 1;
174 }
175 return try out.toOwnedSlice(allocator);
176 }
177
178 pub fn unionHits(allocator: Allocator, left: []const Hit, right: []const Hit) engine_mod.Error![]Hit {
179 if (left.len == 0) return try allocator.dupe(Hit, right);
180 if (right.len == 0) return try allocator.dupe(Hit, left);
181 const out = try allocator.alloc(Hit, left.len + right.len);
182 errdefer allocator.free(out);
183 if (left[left.len - 1].rowid < right[0].rowid) {
184 @memcpy(out[0..left.len], left);
185 @memcpy(out[left.len..], right);
186 return out;
187 }
188 if (right[right.len - 1].rowid < left[0].rowid) {
189 @memcpy(out[0..right.len], right);
190 @memcpy(out[right.len..], left);
191 return out;
192 }
193 var write: usize = 0;
194 var left_index: usize = 0;
195 var right_index: usize = 0;
196 while (left_index < left.len and right_index < right.len) {
197 if (left[left_index].rowid == right[right_index].rowid) {
198 out[write] = .{
199 .rowid = left[left_index].rowid,
200 .score = left[left_index].score + right[right_index].score,
201 };
202 left_index += 1;
203 right_index += 1;
204 } else if (left[left_index].rowid < right[right_index].rowid) {
205 out[write] = left[left_index];
206 left_index += 1;
207 } else {
208 out[write] = right[right_index];
209 right_index += 1;
210 }
211 write += 1;
212 }
213 while (left_index < left.len) : (left_index += 1) {
214 out[write] = left[left_index];
215 write += 1;
216 }
217 while (right_index < right.len) : (right_index += 1) {
218 out[write] = right[right_index];
219 write += 1;
220 }
221 return try allocator.realloc(out, write);
222 }
223
224 pub fn appendBoundedHit(allocator: Allocator, hits: *std.ArrayList(Hit), limit: usize, hit: Hit) engine_mod.Error!void {
225 std.debug.assert(limit != 0);
226 if (hits.items.len < limit) {
227 try hits.append(allocator, hit);
228 siftUpWorstHit(hits.items, hits.items.len - 1);
229 return;
230 }
231 if (!hitLess({}, hit, hits.items[0])) return;
232 hits.items[0] = hit;
233 siftDownWorstHit(hits.items, 0);
234 }
235
236 pub fn appendRankedHit(allocator: Allocator, hits: *std.ArrayList(RankedHit), limit: usize, hit: RankedHit) engine_mod.Error!void {
237 if (limit == 0) {
238 try hits.append(allocator, hit);
239 return;
240 }
241 if (hits.items.len < limit) {
242 try hits.append(allocator, hit);
243 siftUpWorstRankedHit(hits.items, hits.items.len - 1);
244 return;
245 }
246 if (!rankedHitLess({}, hit, hits.items[0])) return;
247 hits.items[0] = hit;
248 siftDownWorstRankedHit(hits.items, 0);
249 }
250
251 pub fn sortAndLimitOwnedHits(allocator: Allocator, hits: []Hit, limit: usize) engine_mod.Error![]Hit {
252 if (limit == 0 or hits.len <= limit) {
253 std.mem.sort(Hit, hits, {}, hitLess);
254 return hits;
255 }
256 if (hitsRanked(hits)) return try allocator.realloc(hits, limit);
257 var heap_len: usize = 0;
258 while (heap_len < limit) : (heap_len += 1) {
259 siftUpWorstHit(hits[0 .. heap_len + 1], heap_len);
260 }
261 var index = limit;
262 while (index < hits.len) : (index += 1) {
263 if (!hitLess({}, hits[index], hits[0])) continue;
264 hits[0] = hits[index];
265 siftDownWorstHit(hits[0..limit], 0);
266 }
267 std.mem.sort(Hit, hits[0..limit], {}, hitLess);
268 return try allocator.realloc(hits, limit);
269 }
270
271 pub fn sortAndScoreBandOwnedHits(allocator: Allocator, hits: []Hit, limit: usize) engine_mod.Error![]Hit {
272 std.mem.sort(Hit, hits, {}, hitLess);
273 if (limit == 0 or hits.len <= limit) return hits;
274 const cutoff = hits[limit - 1].score;
275 var count = limit;
276 while (count < hits.len and hits[count].score >= cutoff) count += 1;
277 return try allocator.realloc(hits, count);
278 }
279
280 pub fn postingHitFrequency(hits: []const Hit, rowid: i64) usize {
281 var start: usize = 0;
282 var end = hits.len;
283 while (start < end) {
284 const middle = start + (end - start) / 2;
285 if (hits[middle].rowid == rowid) return hits[middle].score;
286 if (hits[middle].rowid < rowid) {
287 start = middle + 1;
288 } else {
289 end = middle;
290 }
291 }
292 return 0;
293 }
294
295 pub fn minHitScore(hits: []const Hit) usize {
296 std.debug.assert(hits.len != 0);
297 var score = hits[0].score;
298 for (hits[1..]) |hit| score = @min(score, hit.score);
299 return score;
300 }
301
302 pub fn hitsRanked(hits: []const Hit) bool {
303 if (hits.len < 2) return true;
304 var index: usize = 1;
305 while (index < hits.len) : (index += 1) {
306 if (hitLess({}, hits[index], hits[index - 1])) return false;
307 }
308 return true;
309 }
310
311 pub fn siftUpWorstHit(hits: []Hit, start: usize) void {
312 var index = start;
313 while (index > 0) {
314 const parent = (index - 1) / 2;
315 if (!hitWorse(hits[index], hits[parent])) break;
316 std.mem.swap(Hit, &hits[index], &hits[parent]);
317 index = parent;
318 }
319 }
320
321 pub fn siftDownWorstHit(hits: []Hit, start: usize) void {
322 var index = start;
323 while (true) {
324 const left = index * 2 + 1;
325 if (left >= hits.len) break;
326 const right = left + 1;
327 const child = if (right < hits.len and hitWorse(hits[right], hits[left])) right else left;
328 if (!hitWorse(hits[child], hits[index])) break;
329 std.mem.swap(Hit, &hits[index], &hits[child]);
330 index = child;
331 }
332 }
333
334 pub fn hitWorse(left: Hit, right: Hit) bool {
335 return hitLess({}, right, left);
336 }
337
338 pub fn siftUpWorstRankedHit(hits: []RankedHit, start: usize) void {
339 var index = start;
340 while (index > 0) {
341 const parent = (index - 1) / 2;
342 if (!rankedHitWorse(hits[index], hits[parent])) break;
343 std.mem.swap(RankedHit, &hits[index], &hits[parent]);
344 index = parent;
345 }
346 }
347
348 pub fn siftDownWorstRankedHit(hits: []RankedHit, start: usize) void {
349 var index = start;
350 while (true) {
351 const left = index * 2 + 1;
352 if (left >= hits.len) break;
353 const right = left + 1;
354 const child = if (right < hits.len and rankedHitWorse(hits[right], hits[left])) right else left;
355 if (!rankedHitWorse(hits[child], hits[index])) break;
356 std.mem.swap(RankedHit, &hits[index], &hits[child]);
357 index = child;
358 }
359 }
360
361 pub fn rankedHitWorse(left: RankedHit, right: RankedHit) bool {
362 return rankedHitLess({}, right, left);
363 }
364
365 pub fn sortedUniqueRows(allocator: Allocator, rows: *std.ArrayList(i64)) engine_mod.Error![]i64 {
366 std.mem.sort(i64, rows.items, {}, rowidLess);
367 var write: usize = 0;
368 for (rows.items) |rowid| {
369 if (write != 0 and rows.items[write - 1] == rowid) continue;
370 rows.items[write] = rowid;
371 write += 1;
372 }
373 rows.shrinkRetainingCapacity(write);
374 return try rows.toOwnedSlice(allocator);
375 }
376
377 pub fn rowidLess(_: void, left: i64, right: i64) bool {
378 return left < right;
379 }
380
381 pub fn hitLess(_: void, left: Hit, right: Hit) bool {
382 if (left.score != right.score) return left.score > right.score;
383 return left.rowid < right.rowid;
384 }
385
386 pub fn rankedHitLess(_: void, left: RankedHit, right: RankedHit) bool {
387 if (left.rank != right.rank) return left.rank < right.rank;
388 return left.rowid < right.rowid;
389 }