GCC Code Coverage Report

Directory: src/
File: src/algo/topk.h
Date: 2026-10-23 09:17:32
Exec Total Coverage
Lines: 4 4 100.0%
Functions: 1 1 100.0%
Branches: 10 10 100.0%
Line Branch Exec Source
1 /*
2 * Copyright (c) 2026 Tiger Data, Inc.
3 * Licensed under the PostgreSQL License. See LICENSE for details.
4 *
5 * topk.h - Bounded top-K collection for nearest neighbor results
6 *
7 * Collects candidates that could be in the true top-K based on
8 * RaBitQ error bounds. Two internal structures work together:
9 *
10 * 1. Threshold heap: max-heap of K upper bounds (distance + error).
11 * Tracks the Kth-smallest upper bound seen so far. A candidate
12 * is pruned when its lower bound >= this threshold (there are
13 * already K candidates that are definitely closer).
14 *
15 * 2. Candidate buffer: growable array of all entries that passed
16 * the threshold check. May contain more than K entries because
17 * error intervals can overlap near the Kth position.
18 *
19 * After the scan, the buffer contains the full rerank set. The
20 * caller fetches full-precision vectors for overlapping candidates,
21 * reranks, and takes the true top-K.
22 *
23 * Usage:
24 * VsTopK topk;
25 * vs_topk_init(&topk, k);
26 *
27 * vs_topk_insert(&topk, distance, error, id);
28 * // ... more inserts ...
29 *
30 * VsTopKEntry *results = vs_alloc(topk.cand_count * sizeof(...));
31 * uint32_t count;
32 * vs_topk_extract_sorted(&topk, results, &count);
33 * vs_topk_cleanup(&topk);
34 */
35
36 #ifndef VS_TOPK_H
37 #define VS_TOPK_H
38
39 #include <math.h>
40 #include <stdint.h>
41
42 #include "core/memory.h"
43 #include "core/types.h"
44
45 /* ----------------------------------------------------------------
46 * Top-K entry
47 * ---------------------------------------------------------------- */
48 typedef struct VsTopKEntry
49 {
50 Distance distance; /* estimated distance */
51 Distance error; /* symmetric error margin (>= 0) */
52 uint64_t id; /* encoded TID or generic identifier */
53 uint32_t src; /* diagnostic: source rank stamped at insert */
54 } VsTopKEntry;
55
56 /* ----------------------------------------------------------------
57 * Top-K collection
58 *
59 * Binary max-heap array for threshold tracking, in both the standalone
60 * and PG builds.
61 * ---------------------------------------------------------------- */
62 typedef struct VsTopK
63 {
64 Distance *ub_heap; /* binary max-heap of K upper bounds */
65 uint64_t *ub_ids; /* parallel array: ID for each heap entry */
66 uint32_t ub_count; /* entries in threshold heap (<= k) */
67 uint32_t k; /* target K */
68 uint32_t k_capacity; /* allocated ub_heap/ub_ids capacity */
69 VsTopKEntry *candidates; /* growable candidate buffer */
70 uint32_t cand_count; /* buffered candidates */
71 uint32_t cand_capacity; /* allocated capacity */
72 uint32_t cur_src; /* diagnostic: rank stamped onto inserts */
73 VsMemCtx memctx; /* owning context for all allocations */
74
75 /* Id -> heap position map for vs_topk_insert's dedup lookup, built
76 * on the first deduping insert. */
77 uint64_t *map_ids; /* open-addressing keys */
78 uint32_t *map_pos; /* heap position per slot, UINT32_MAX if free */
79 uint32_t *ub_slot; /* map slot per heap position */
80 uint32_t map_mask; /* map slots - 1 */
81 uint32_t map_shift; /* 64 - log2(map slots) */
82 } VsTopK;
83
84 /* ----------------------------------------------------------------
85 * API
86 * ---------------------------------------------------------------- */
87
88 /*
89 * Initialize a top-K collection. Creates a child memory context
90 * under the current context for internal buffers. Caller must be
91 * in a context with the desired lifetime.
92 * Use vs_topk_cleanup() to free.
93 */
94 void vs_topk_init(VsTopK *topk, uint32_t k);
95
96 /* Free internal buffers (does not free the VsTopK struct itself). */
97 void vs_topk_cleanup(VsTopK *topk);
98
99 /*
100 * Create a heap-allocated top-K collection.
101 * Returns NULL on allocation failure.
102 */
103 VsTopK *vs_topk_create(uint32_t k);
104
105 /* Free a heap-allocated top-K collection. */
106 void vs_topk_destroy(VsTopK *topk);
107
108 /* Reset to empty state (reuses existing storage). */
109 void vs_topk_reset(VsTopK *topk);
110
111 /*
112 * Reset to empty state AND change k. Reuses the existing memory
113 * context (preserves the per-topk arena) but re-allocates ub_heap /
114 * ub_ids / candidates within it. Cheap compared to a full init +
115 * cleanup pair, because the memctx itself isn't created or destroyed.
116 *
117 * Use this when the same VsTopK is reused across calls that may
118 * want different k values (e.g. beam-search keeps beam_width
119 * candidates at intermediate levels, nprobe at the last).
120 */
121 void vs_topk_reset_to_k(VsTopK *topk, uint32_t k);
122
123 /*
124 * Same as vs_topk_insert but skips the per-insert dedup lookup.
125 * Use ONLY when the caller guarantees all ids are unique. The cluster
126 * scan (where SOAR / boundary replicas can collide) must keep using
127 * vs_topk_insert; centroid beam search and similar code that
128 * inserts each candidate exactly once should use this fast path.
129 * It keeps a dedup map that vs_topk_insert built up to date, so the
130 * two can be mixed on one top-K.
131 */
132 void vs_topk_insert_unique(
133 VsTopK *topk, Distance distance, Distance error, uint64_t id);
134
135 /*
136 * Insert a candidate. Pruned if lower_bound >= threshold.
137 * Otherwise updates the threshold heap and appends to the
138 * candidate buffer.
139 */
140 void
141 vs_topk_insert(VsTopK *topk, Distance distance, Distance error, uint64_t id);
142
143 /*
144 * Current pruning threshold: the Kth-smallest upper bound
145 * (distance + error) seen so far. Returns INFINITY if fewer
146 * than K upper bounds have been recorded.
147 */
148 static inline Distance
149 56693581 vs_topk_threshold(const VsTopK *topk)
150 {
151
10/10
✓ Branch 0 taken 3569842 times.
✓ Branch 1 taken 11155372 times.
✓ Branch 2 taken 6600 times.
✓ Branch 3 taken 5221 times.
✓ Branch 4 taken 29338 times.
✓ Branch 5 taken 9768 times.
✓ Branch 6 taken 29228865 times.
✓ Branch 7 taken 12684186 times.
✓ Branch 8 taken 2992 times.
✓ Branch 9 taken 1397 times.
 56693581 if (topk->ub_count < topk->k)
152 2977574 return INFINITY;
153 40906653 return topk->ub_heap[0];
154 }
155
156 /*
157 * Extract candidates sorted by distance ascending. Filters out
158 * stale entries whose lower bound now exceeds the final threshold.
159 *
160 * results[] must have space for topk->cand_count entries (upper
161 * bound; actual count returned via *count_out may be smaller).
162 *
163 * Does not reset the collection — call vs_topk_reset() or
164 * vs_topk_cleanup() when done with the results.
165 */
166 void vs_topk_extract_sorted_capped(
167 VsTopK *topk, VsTopKEntry *results, uint32_t *count_out, uint32_t cap);
168 void vs_topk_extract_sorted(
169 VsTopK *topk, VsTopKEntry *results, uint32_t *count_out);
170
171 /* As above, but skips the duplicate-id pass; ids must be unique. */
172 void vs_topk_extract_sorted_unique(
173 VsTopK *topk, VsTopKEntry *results, uint32_t *count_out);
174
175 #endif /* VS_TOPK_H */
176