| 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.c - Top-K collection | ||
| 6 | * | ||
| 7 | * Threshold heap: max-heap of K Distance values (upper bounds). | ||
| 8 | * Candidate buffer: growable array of all entries that passed | ||
| 9 | * the threshold check. | ||
| 10 | * Id map: heap position of each id, for vs_topk_insert's dedup lookup. | ||
| 11 | * | ||
| 12 | * Used by both standalone and PG builds. | ||
| 13 | */ | ||
| 14 | |||
| 15 | #include <stdlib.h> | ||
| 16 | #include <string.h> | ||
| 17 | |||
| 18 | #include "algo/topk.h" | ||
| 19 | #include "core/idset.h" | ||
| 20 | #include "core/memory.h" | ||
| 21 | #include "core/platform.h" | ||
| 22 | |||
| 23 | #define VS_TOPK_INITIAL_CAP_MIN 32 | ||
| 24 | |||
| 25 | /* ---------------------------------------------------------------- | ||
| 26 | * Id -> heap position map | ||
| 27 | * | ||
| 28 | * Linear probing over a power-of-two slot array sized for at most half | ||
| 29 | * load at k_capacity. Deletion shifts later entries back into the hole | ||
| 30 | * rather than leaving tombstones, so probe chains stay short however | ||
| 31 | * many evictions a scan makes. | ||
| 32 | * ---------------------------------------------------------------- */ | ||
| 33 | |||
| 34 | #define VS_TOPK_MAP_EMPTY UINT32_MAX | ||
| 35 | |||
| 36 | static inline uint32_t | ||
| 37 | 723689 | map_home(const VsTopK *topk, uint64_t id) | |
| 38 | { | ||
| 39 | 723689 | return (uint32_t)((id * VS_HASH_GOLDEN_GAMMA) >> topk->map_shift); | |
| 40 | } | ||
| 41 | |||
| 42 | /* Slot holding `id`, or VS_TOPK_MAP_EMPTY if it is not in the heap. */ | ||
| 43 | static inline uint32_t | ||
| 44 | 646261 | map_find(const VsTopK *topk, uint64_t id) | |
| 45 | { | ||
| 46 | 646261 | uint32_t slot = map_home(topk, id); | |
| 47 |
2/2✓ Branch 0 taken 300738 times.
✓ Branch 1 taken 640406 times.
|
941144 | while (topk->map_pos[slot] != VS_TOPK_MAP_EMPTY) |
| 48 | { | ||
| 49 |
2/2✓ Branch 0 taken 5855 times.
✓ Branch 1 taken 294883 times.
|
300738 | if (topk->map_ids[slot] == id) |
| 50 | 5855 | return slot; | |
| 51 | 294883 | slot = (slot + 1) & topk->map_mask; | |
| 52 | } | ||
| 53 | 613558 | return VS_TOPK_MAP_EMPTY; | |
| 54 | } | ||
| 55 | |||
| 56 | /* Record that heap position `pos` holds `id` (absent from the map). */ | ||
| 57 | static inline void | ||
| 58 | 66697 | map_add(VsTopK *topk, uint64_t id, uint32_t pos) | |
| 59 | { | ||
| 60 | 66697 | uint32_t slot = map_home(topk, id); | |
| 61 |
2/2✓ Branch 0 taken 6694 times.
✓ Branch 1 taken 66697 times.
|
73391 | while (topk->map_pos[slot] != VS_TOPK_MAP_EMPTY) |
| 62 | 6694 | slot = (slot + 1) & topk->map_mask; | |
| 63 | 66697 | topk->map_ids[slot] = id; | |
| 64 | 66697 | topk->map_pos[slot] = pos; | |
| 65 | 66697 | topk->ub_slot[pos] = slot; | |
| 66 | 66697 | } | |
| 67 | |||
| 68 | static void | ||
| 69 | 45126 | map_remove(VsTopK *topk, uint32_t slot) | |
| 70 | { | ||
| 71 | 45126 | uint32_t mask = topk->map_mask; | |
| 72 | 45126 | uint32_t hole = slot; | |
| 73 | |||
| 74 |
2/2✓ Branch 0 taken 10731 times.
✓ Branch 1 taken 45126 times.
|
55857 | for (uint32_t j = (hole + 1) & mask; topk->map_pos[j] != VS_TOPK_MAP_EMPTY; |
| 75 | 10731 | j = (j + 1) & mask) | |
| 76 | { | ||
| 77 | /* Move j into the hole unless that would put it before its | ||
| 78 | * home slot, where lookups would no longer reach it. */ | ||
| 79 | 10731 | uint32_t home = map_home(topk, topk->map_ids[j]); | |
| 80 |
2/2✓ Branch 0 taken 2373 times.
✓ Branch 1 taken 8358 times.
|
10731 | if (((j - home) & mask) >= ((j - hole) & mask)) |
| 81 | { | ||
| 82 | 2373 | topk->map_ids[hole] = topk->map_ids[j]; | |
| 83 | 2373 | topk->map_pos[hole] = topk->map_pos[j]; | |
| 84 | 2373 | topk->ub_slot[topk->map_pos[hole]] = hole; | |
| 85 | 2373 | hole = j; | |
| 86 | } | ||
| 87 | } | ||
| 88 | 45126 | topk->map_pos[hole] = VS_TOPK_MAP_EMPTY; | |
| 89 | 45126 | } | |
| 90 | |||
| 91 | /* Allocate the map on the first deduping insert, indexing whatever | ||
| 92 | * unique inserts already put in the heap. */ | ||
| 93 | static void | ||
| 94 | 471 | map_build(VsTopK *topk) | |
| 95 | { | ||
| 96 | 471 | uint32_t bits = 1; | |
| 97 |
2/2✓ Branch 0 taken 2058 times.
✓ Branch 1 taken 471 times.
|
2529 | while ((1ull << bits) < 2ull * topk->k_capacity) |
| 98 | 2058 | bits++; | |
| 99 | |||
| 100 | 471 | uint32_t nslots = 1u << bits; | |
| 101 | 471 | topk->map_ids = vs_memctx_alloc(topk->memctx, nslots * sizeof(uint64_t)); | |
| 102 | 471 | topk->map_pos = vs_memctx_alloc(topk->memctx, nslots * sizeof(uint32_t)); | |
| 103 | 738 | topk->ub_slot = | |
| 104 | 471 | vs_memctx_alloc(topk->memctx, topk->k_capacity * sizeof(uint32_t)); | |
| 105 | 471 | topk->map_mask = nslots - 1; | |
| 106 | 471 | topk->map_shift = 64 - bits; | |
| 107 | |||
| 108 | 471 | memset(topk->map_pos, 0xFF, nslots * sizeof(uint32_t)); | |
| 109 |
2/2✓ Branch 0 taken 64 times.
✓ Branch 1 taken 471 times.
|
535 | for (uint32_t i = 0; i < topk->ub_count; i++) |
| 110 | 64 | map_add(topk, topk->ub_ids[i], i); | |
| 111 | 471 | } | |
| 112 | |||
| 113 | /* ---------------------------------------------------------------- | ||
| 114 | * Threshold heap (max-heap of Distance values) | ||
| 115 | * | ||
| 116 | * The sifts take the heap's arrays as locals, so stores through them | ||
| 117 | * don't force reloads of the pointers from VsTopK. With `track`, every | ||
| 118 | * move also updates the id map. | ||
| 119 | * ---------------------------------------------------------------- */ | ||
| 120 | |||
| 121 | typedef struct HeapArrays | ||
| 122 | { | ||
| 123 | Distance *ub; | ||
| 124 | uint64_t *ids; | ||
| 125 | uint32_t *slot; /* ub_slot, used only when tracking */ | ||
| 126 | uint32_t *pos; /* map_pos, used only when tracking */ | ||
| 127 | } HeapArrays; | ||
| 128 | |||
| 129 | static inline HeapArrays | ||
| 130 | 27695486 | heap_arrays(const VsTopK *topk) | |
| 131 | { | ||
| 132 | 33314050 | return (HeapArrays){ | |
| 133 | 27695486 | .ub = topk->ub_heap, | |
| 134 | 27695486 | .ids = topk->ub_ids, | |
| 135 | 27695486 | .slot = topk->ub_slot, | |
| 136 | 27695486 | .pos = topk->map_pos, | |
| 137 | }; | ||
| 138 | } | ||
| 139 | |||
| 140 | static inline void | ||
| 141 | 54575710 | ub_swap(HeapArrays h, uint32_t a, uint32_t b, bool track) | |
| 142 | { | ||
| 143 | 54575710 | Distance tmp = h.ub[a]; | |
| 144 | 54575710 | h.ub[a] = h.ub[b]; | |
| 145 | 54575710 | h.ub[b] = tmp; | |
| 146 | 54575710 | uint64_t tid = h.ids[a]; | |
| 147 | 54575710 | h.ids[a] = h.ids[b]; | |
| 148 | 54575710 | h.ids[b] = tid; | |
| 149 | |||
| 150 |
2/2✓ Branch 0 taken 99158 times.
✓ Branch 1 taken 54476552 times.
|
54575710 | if (track) |
| 151 | { | ||
| 152 | 99158 | uint32_t sa = h.slot[a]; | |
| 153 | 99158 | uint32_t sb = h.slot[b]; | |
| 154 | 99158 | h.slot[a] = sb; | |
| 155 | 99158 | h.slot[b] = sa; | |
| 156 | 99158 | h.pos[sb] = a; | |
| 157 | 99158 | h.pos[sa] = b; | |
| 158 | } | ||
| 159 | 54575710 | } | |
| 160 | |||
| 161 | static inline void | ||
| 162 | 15502633 | ub_sift_up(HeapArrays h, uint32_t i, bool track) | |
| 163 | { | ||
| 164 |
2/2✓ Branch 0 taken 28716873 times.
✓ Branch 1 taken 3095400 times.
|
31812273 | while (i > 0) |
| 165 | { | ||
| 166 | 28716873 | uint32_t parent = (i - 1) / 2; | |
| 167 |
2/2✓ Branch 0 taken 15713098 times.
✓ Branch 1 taken 13003775 times.
|
28716873 | if (h.ub[i] <= h.ub[parent]) |
| 168 | 2175232 | break; | |
| 169 | 16309640 | ub_swap(h, i, parent, track); | |
| 170 | 16309640 | i = parent; | |
| 171 | } | ||
| 172 | 15502633 | } | |
| 173 | |||
| 174 | static inline void | ||
| 175 | 12192853 | ub_sift_down(HeapArrays h, uint32_t count, uint32_t i, bool track) | |
| 176 | { | ||
| 177 | 68172873 | for (;;) | |
| 178 | 38266070 | { | |
| 179 | 50458923 | uint32_t left = 2 * i + 1; | |
| 180 | 50458923 | uint32_t right = 2 * i + 2; | |
| 181 | 50458923 | uint32_t largest = i; | |
| 182 | |||
| 183 |
4/4✓ Branch 0 taken 15338694 times.
✓ Branch 1 taken 35120229 times.
✓ Branch 2 taken 15033516 times.
✓ Branch 3 taken 28939665 times.
|
50458923 | if (left < count && h.ub[left] > h.ub[largest]) |
| 184 | 46775526 | largest = left; | |
| 185 |
4/4✓ Branch 0 taken 15724859 times.
✓ Branch 1 taken 34734064 times.
✓ Branch 2 taken 22619992 times.
✓ Branch 3 taken 20641148 times.
|
50458923 | if (right < count && h.ub[right] > h.ub[largest]) |
| 186 | 43183737 | largest = right; | |
| 187 | |||
| 188 |
2/2✓ Branch 0 taken 32202190 times.
✓ Branch 1 taken 18256733 times.
|
50458923 | if (largest == i) |
| 189 | 2808120 | break; | |
| 190 | |||
| 191 | 38266070 | ub_swap(h, i, largest, track); | |
| 192 | 38266070 | i = largest; | |
| 193 | } | ||
| 194 | 12192853 | } | |
| 195 | |||
| 196 | /* Add an id that isn't in the heap. It is appended while the heap holds | ||
| 197 | * fewer than k entries and otherwise replaces the root if ub beats it. */ | ||
| 198 | __attribute__((always_inline)) static inline void | ||
| 199 | 22089994 | ub_insert(VsTopK *topk, Distance ub, uint64_t id, bool track) | |
| 200 | { | ||
| 201 |
6/6✓ Branch 0 taken 13342 times.
✓ Branch 1 taken 600216 times.
✓ Branch 2 taken 2797014 times.
✓ Branch 3 taken 2767321 times.
✓ Branch 4 taken 88 times.
✓ Branch 5 taken 340 times.
|
28268315 | if (topk->ub_count < topk->k) |
| 202 | { | ||
| 203 | 15502633 | uint32_t i = topk->ub_count++; | |
| 204 | 15502633 | topk->ub_heap[i] = ub; | |
| 205 | 15502633 | topk->ub_ids[i] = id; | |
| 206 |
3/6✓ Branch 0 taken 13342 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2797014 times.
✓ Branch 4 taken 88 times.
✗ Branch 5 not taken.
|
15502633 | if (track) |
| 207 | 21507 | map_add(topk, id, i); | |
| 208 | 15502633 | ub_sift_up(heap_arrays(topk), i, track); | |
| 209 | } | ||
| 210 |
5/6✓ Branch 0 taken 44968 times.
✓ Branch 1 taken 574019 times.
✓ Branch 2 taken 12146355 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 158 times.
✓ Branch 5 taken 182 times.
|
12765682 | else if (ub < topk->ub_heap[0]) |
| 211 | { | ||
| 212 |
3/6✓ Branch 0 taken 39367 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2767321 times.
✓ Branch 4 taken 158 times.
✗ Branch 5 not taken.
|
12191481 | if (track) |
| 213 | 45126 | map_remove(topk, topk->ub_slot[0]); | |
| 214 | 12191481 | topk->ub_heap[0] = ub; | |
| 215 | 12191481 | topk->ub_ids[0] = id; | |
| 216 |
3/6✓ Branch 0 taken 39367 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2767321 times.
✓ Branch 4 taken 158 times.
✗ Branch 5 not taken.
|
12191481 | if (track) |
| 217 | 45126 | map_add(topk, id, 0); | |
| 218 | 12191481 | ub_sift_down(heap_arrays(topk), topk->ub_count, 0, track); | |
| 219 | } | ||
| 220 | 6178321 | } | |
| 221 | |||
| 222 | /* ---------------------------------------------------------------- | ||
| 223 | * Comparator for qsort (ascending by distance) | ||
| 224 | * ---------------------------------------------------------------- */ | ||
| 225 | |||
| 226 | static int | ||
| 227 | 68194569 | cmp_by_distance(const void *a, const void *b) | |
| 228 | { | ||
| 229 | 68194569 | const VsTopKEntry *ea = (const VsTopKEntry *)a; | |
| 230 | 68194569 | const VsTopKEntry *eb = (const VsTopKEntry *)b; | |
| 231 |
2/2✓ Branch 0 taken 35662697 times.
✓ Branch 1 taken 32531872 times.
|
68194569 | if (ea->distance < eb->distance) |
| 232 | 5091979 | return -1; | |
| 233 |
2/2✓ Branch 0 taken 35409023 times.
✓ Branch 1 taken 30060 times.
|
35439083 | if (ea->distance > eb->distance) |
| 234 | 35409023 | return 1; | |
| 235 | 28008 | return 0; | |
| 236 | } | ||
| 237 | |||
| 238 | /* ---------------------------------------------------------------- | ||
| 239 | * Init / cleanup / create / destroy | ||
| 240 | * ---------------------------------------------------------------- */ | ||
| 241 | |||
| 242 | void | ||
| 243 | 36722 | vs_topk_init(VsTopK *topk, uint32_t k) | |
| 244 | { | ||
| 245 | 36722 | topk->memctx = vs_memctx_create(NULL, "topk"); | |
| 246 | 36722 | topk->k = k; | |
| 247 | 36722 | topk->k_capacity = k; | |
| 248 | 36722 | topk->ub_heap = vs_memctx_alloc(topk->memctx, k * sizeof(Distance)); | |
| 249 | 36722 | topk->ub_ids = vs_memctx_alloc(topk->memctx, k * sizeof(uint64_t)); | |
| 250 | 36722 | topk->ub_count = 0; | |
| 251 | |||
| 252 | 36722 | uint32_t cap = k * 2; | |
| 253 |
2/2✓ Branch 0 taken 334 times.
✓ Branch 1 taken 258 times.
|
36722 | if (cap < VS_TOPK_INITIAL_CAP_MIN) |
| 254 | 334 | cap = VS_TOPK_INITIAL_CAP_MIN; | |
| 255 | 72852 | topk->candidates = | |
| 256 | 36722 | vs_memctx_alloc(topk->memctx, cap * sizeof(VsTopKEntry)); | |
| 257 | 36722 | topk->cand_count = 0; | |
| 258 | 36722 | topk->cand_capacity = cap; | |
| 259 | /* Stamped onto every insert; only the query scan ever sets it. */ | ||
| 260 | 36722 | topk->cur_src = 0; | |
| 261 | |||
| 262 | 36722 | topk->map_ids = NULL; | |
| 263 | 36722 | topk->map_pos = NULL; | |
| 264 | 36722 | topk->ub_slot = NULL; | |
| 265 | 36722 | topk->map_mask = 0; | |
| 266 | 36722 | topk->map_shift = 0; | |
| 267 | 36722 | } | |
| 268 | |||
| 269 | void | ||
| 270 | 36718 | vs_topk_cleanup(VsTopK *topk) | |
| 271 | { | ||
| 272 |
2/2✓ Branch 0 taken 36126 times.
✓ Branch 1 taken 592 times.
|
36718 | if (topk == NULL) |
| 273 | ✗ | return; | |
| 274 |
1/2✓ Branch 0 taken 36718 times.
✗ Branch 1 not taken.
|
36718 | if (topk->memctx != NULL) |
| 275 | { | ||
| 276 | 36718 | vs_memctx_delete(topk->memctx); | |
| 277 | 36718 | topk->memctx = NULL; | |
| 278 | } | ||
| 279 | 36718 | topk->ub_heap = NULL; | |
| 280 | 36718 | topk->ub_ids = NULL; | |
| 281 | 36718 | topk->candidates = NULL; | |
| 282 | 36718 | topk->map_ids = NULL; | |
| 283 | 36718 | topk->map_pos = NULL; | |
| 284 | 36718 | topk->ub_slot = NULL; | |
| 285 | } | ||
| 286 | |||
| 287 | VsTopK * | ||
| 288 | 6 | vs_topk_create(uint32_t k) | |
| 289 | { | ||
| 290 | /* The handle comes from the caller's context rather than from the arena | ||
| 291 | * it owns. vs_topk_reset_to_k resets that arena when k outgrows it, | ||
| 292 | * which would reclaim the object being reset through. */ | ||
| 293 | 6 | VsTopK *topk = vs_alloc(sizeof(VsTopK)); | |
| 294 | 6 | vs_topk_init(topk, k); | |
| 295 | 6 | return topk; | |
| 296 | } | ||
| 297 | |||
| 298 | void | ||
| 299 | 6 | vs_topk_destroy(VsTopK *topk) | |
| 300 | { | ||
| 301 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
|
6 | if (topk == NULL) |
| 302 | ✗ | return; | |
| 303 | 6 | vs_topk_cleanup(topk); | |
| 304 | 6 | vs_free(topk); | |
| 305 | } | ||
| 306 | |||
| 307 | void | ||
| 308 | 4 | vs_topk_reset(VsTopK *topk) | |
| 309 | { | ||
| 310 | 4 | vs_topk_reset_to_k(topk, topk->k); | |
| 311 | 4 | } | |
| 312 | |||
| 313 | void | ||
| 314 | 909738 | vs_topk_reset_to_k(VsTopK *topk, uint32_t k) | |
| 315 | { | ||
| 316 | /* O(1) reset: buffers are retained across resets (grow-only), so a | ||
| 317 | * reused top-K reaches a steady state with zero allocator traffic. | ||
| 318 | * Only a k above the allocated capacity re-allocates — reset the | ||
| 319 | * arena first so the outgrown blocks are reclaimed rather than | ||
| 320 | * leaked into the context. The candidate buffer keeps its | ||
| 321 | * high-water capacity, avoiding the doubling-regrowth copies that | ||
| 322 | * a fresh buffer paid on every use. */ | ||
| 323 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 909722 times.
|
909738 | if (vs_unlikely(k > topk->k_capacity)) |
| 324 | { | ||
| 325 | 16 | uint32_t cand_cap = topk->cand_capacity; | |
| 326 | 16 | uint32_t min_cap = k * 2; | |
| 327 | |||
| 328 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 14 times.
|
16 | if (cand_cap < min_cap) |
| 329 | 2 | cand_cap = min_cap; | |
| 330 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
|
16 | if (cand_cap < VS_TOPK_INITIAL_CAP_MIN) |
| 331 | ✗ | cand_cap = VS_TOPK_INITIAL_CAP_MIN; | |
| 332 | |||
| 333 | 16 | vs_memctx_reset(topk->memctx); | |
| 334 | 16 | topk->ub_heap = vs_memctx_alloc(topk->memctx, k * sizeof(Distance)); | |
| 335 | 16 | topk->ub_ids = vs_memctx_alloc(topk->memctx, k * sizeof(uint64_t)); | |
| 336 | 16 | topk->candidates = | |
| 337 | 16 | vs_memctx_alloc(topk->memctx, cand_cap * sizeof(VsTopKEntry)); | |
| 338 | 16 | topk->cand_capacity = cand_cap; | |
| 339 | 16 | topk->k_capacity = k; | |
| 340 | 16 | topk->map_ids = NULL; | |
| 341 | 16 | topk->map_pos = NULL; | |
| 342 | 16 | topk->ub_slot = NULL; | |
| 343 | } | ||
| 344 |
2/2✓ Branch 0 taken 7411 times.
✓ Branch 1 taken 902311 times.
|
909722 | else if (topk->map_pos != NULL) |
| 345 | { | ||
| 346 | /* Clear only the occupied slots. */ | ||
| 347 |
2/2✓ Branch 0 taken 10446 times.
✓ Branch 1 taken 7411 times.
|
17857 | for (uint32_t i = 0; i < topk->ub_count; i++) |
| 348 | 10446 | topk->map_pos[topk->ub_slot[i]] = VS_TOPK_MAP_EMPTY; | |
| 349 | } | ||
| 350 | |||
| 351 | 909738 | topk->k = k; | |
| 352 | 909738 | topk->ub_count = 0; | |
| 353 | 909738 | topk->cand_count = 0; | |
| 354 | 909738 | } | |
| 355 | |||
| 356 | void | ||
| 357 | 54171956 | vs_topk_insert_unique( | |
| 358 | VsTopK *topk, Distance distance, Distance error, uint64_t id) | ||
| 359 | { | ||
| 360 | 54171956 | Distance lb = distance - error; | |
| 361 | 54171956 | Distance ub = distance + error; | |
| 362 | |||
| 363 | /* Prune: lower bound exceeds threshold. Written NaN-safe: a NaN | ||
| 364 | * estimate (e.g. a zero vector normalized for cosine) fails every | ||
| 365 | * comparison, so the naive lb >= threshold guard would admit it -- | ||
| 366 | * and one NaN upper bound in the heap corrupts the threshold for | ||
| 367 | * the rest of the scan. */ | ||
| 368 |
4/4✓ Branch 0 taken 29227288 times.
✓ Branch 1 taken 19379905 times.
✓ Branch 2 taken 27627909 times.
✓ Branch 3 taken 19848254 times.
|
83399244 | if (!(lb < vs_topk_threshold(topk))) |
| 369 | 6695793 | return; | |
| 370 | |||
| 371 | /* No dedup lookup: caller guarantees ids are unique. A map that | ||
| 372 | * vs_topk_insert built is kept current. */ | ||
| 373 |
2/2✓ Branch 0 taken 27627481 times.
✓ Branch 1 taken 428 times.
|
27627909 | if (topk->map_pos == NULL) |
| 374 |
2/2✓ Branch 0 taken 12684112 times.
✓ Branch 1 taken 9379034 times.
|
22063146 | ub_insert(topk, ub, id, false); |
| 375 | else | ||
| 376 | ✗ | ub_insert(topk, ub, id, true); | |
| 377 | |||
| 378 | /* Grow candidate buffer if needed (old buffer freed with memctx) */ | ||
| 379 |
2/2✓ Branch 0 taken 36 times.
✓ Branch 1 taken 27627873 times.
|
27627909 | if (topk->cand_count == topk->cand_capacity) |
| 380 | { | ||
| 381 | 36 | uint32_t new_cap = topk->cand_capacity * 2; | |
| 382 | 24 | VsTopKEntry *new_buf = | |
| 383 | 36 | vs_memctx_alloc(topk->memctx, new_cap * sizeof(VsTopKEntry)); | |
| 384 | 36 | memcpy(new_buf, | |
| 385 | 36 | topk->candidates, | |
| 386 | 36 | topk->cand_count * sizeof(VsTopKEntry)); | |
| 387 | 36 | topk->candidates = new_buf; | |
| 388 | 36 | topk->cand_capacity = new_cap; | |
| 389 | } | ||
| 390 | |||
| 391 | 27627909 | topk->candidates[topk->cand_count++] = (VsTopKEntry){ | |
| 392 | .distance = distance, | ||
| 393 | .error = error, | ||
| 394 | .id = id, | ||
| 395 | 27627909 | .src = topk->cur_src, | |
| 396 | }; | ||
| 397 | } | ||
| 398 | |||
| 399 | /* ---------------------------------------------------------------- | ||
| 400 | * Insert | ||
| 401 | * ---------------------------------------------------------------- */ | ||
| 402 | |||
| 403 | void | ||
| 404 | 977227 | vs_topk_insert(VsTopK *topk, Distance distance, Distance error, uint64_t id) | |
| 405 | { | ||
| 406 | 977227 | Distance lb = distance - error; | |
| 407 | 977227 | Distance ub = distance + error; | |
| 408 | |||
| 409 | /* Prune: lower bound exceeds threshold. Written NaN-safe: a NaN | ||
| 410 | * estimate (e.g. a zero vector normalized for cosine) fails every | ||
| 411 | * comparison, so the naive lb >= threshold guard would admit it -- | ||
| 412 | * and one NaN upper bound in the heap corrupts the threshold for | ||
| 413 | * the rest of the scan. */ | ||
| 414 |
4/4✓ Branch 0 taken 19195 times.
✓ Branch 1 taken 339080 times.
✓ Branch 2 taken 646261 times.
✓ Branch 3 taken 208 times.
|
996422 | if (!(lb < vs_topk_threshold(topk))) |
| 415 | 330758 | return; | |
| 416 | |||
| 417 | /* Dedup: if this ID is already in the heap, keep only the better | ||
| 418 | * (smaller) upper bound. This prevents duplicates (e.g. SOAR | ||
| 419 | * replicas) from occupying multiple heap slots and inflating the | ||
| 420 | * threshold. */ | ||
| 421 |
2/2✓ Branch 0 taken 471 times.
✓ Branch 1 taken 645790 times.
|
646261 | if (vs_unlikely(topk->map_pos == NULL)) |
| 422 | 471 | map_build(topk); | |
| 423 | |||
| 424 | 646261 | uint32_t slot = map_find(topk, id); | |
| 425 |
2/2✓ Branch 0 taken 5855 times.
✓ Branch 1 taken 640406 times.
|
646261 | if (slot != VS_TOPK_MAP_EMPTY) |
| 426 | { | ||
| 427 | 5855 | uint32_t i = topk->map_pos[slot]; | |
| 428 |
2/2✓ Branch 0 taken 1372 times.
✓ Branch 1 taken 4483 times.
|
5855 | if (ub < topk->ub_heap[i]) |
| 429 | { | ||
| 430 | 1372 | topk->ub_heap[i] = ub; | |
| 431 | 1372 | ub_sift_down(heap_arrays(topk), topk->ub_count, i, true); | |
| 432 | } | ||
| 433 | 5855 | goto append; | |
| 434 | } | ||
| 435 | |||
| 436 |
2/2✓ Branch 0 taken 8077 times.
✓ Branch 1 taken 18771 times.
|
26848 | ub_insert(topk, ub, id, true); |
| 437 | |||
| 438 | 646261 | append: | |
| 439 | /* Grow candidate buffer if needed (old buffer freed with memctx) */ | ||
| 440 |
2/2✓ Branch 0 taken 613 times.
✓ Branch 1 taken 645648 times.
|
646261 | if (topk->cand_count == topk->cand_capacity) |
| 441 | { | ||
| 442 | 613 | uint32_t new_cap = topk->cand_capacity * 2; | |
| 443 | 237 | VsTopKEntry *new_buf = | |
| 444 | 613 | vs_memctx_alloc(topk->memctx, new_cap * sizeof(VsTopKEntry)); | |
| 445 | 613 | memcpy(new_buf, | |
| 446 | 613 | topk->candidates, | |
| 447 | 613 | topk->cand_count * sizeof(VsTopKEntry)); | |
| 448 | 613 | topk->candidates = new_buf; | |
| 449 | 613 | topk->cand_capacity = new_cap; | |
| 450 | } | ||
| 451 | |||
| 452 | /* Append to candidate buffer */ | ||
| 453 | 646261 | topk->candidates[topk->cand_count++] = (VsTopKEntry){ | |
| 454 | .distance = distance, | ||
| 455 | .error = error, | ||
| 456 | .id = id, | ||
| 457 | 646261 | .src = topk->cur_src, | |
| 458 | }; | ||
| 459 | } | ||
| 460 | |||
| 461 | /* | ||
| 462 | * Deduplicate a distance-sorted candidate prefix in place, keeping the | ||
| 463 | * first (best-distance) entry per ID, tracked in an id set. A rescan | ||
| 464 | * of prior survivors per entry would go quadratic in the survivor | ||
| 465 | * count; the set keeps this pass linear. Returns the deduplicated | ||
| 466 | * count, additionally truncated to `cap` unique entries when cap > 0. | ||
| 467 | */ | ||
| 468 | static uint32_t | ||
| 469 | 7832 | dedup_sorted_prefix(VsTopKEntry *results, uint32_t n, uint32_t cap) | |
| 470 | { | ||
| 471 |
2/2✓ Branch 0 taken 3856 times.
✓ Branch 1 taken 3976 times.
|
7832 | if (n <= 1) |
| 472 | 3580 | return n; | |
| 473 | |||
| 474 | 276 | VsIdSet seen; | |
| 475 | 4246 | vs_idset_init(&seen, n); | |
| 476 | |||
| 477 | 4246 | uint32_t w = 0; | |
| 478 |
3/3✓ Branch 0 taken 90934 times.
✓ Branch 1 taken 12934 times.
✓ Branch 2 taken 221 times.
|
104089 | for (uint32_t r = 0; r < n; r++) |
| 479 | { | ||
| 480 |
2/2✓ Branch 1 taken 5912 times.
✓ Branch 2 taken 97658 times.
|
103570 | if (!vs_idset_test_add(&seen, results[r].id)) |
| 481 | 5912 | continue; | |
| 482 | 97658 | results[w++] = results[r]; | |
| 483 |
4/4✓ Branch 0 taken 87433 times.
✓ Branch 1 taken 10225 times.
✓ Branch 2 taken 3672 times.
✓ Branch 3 taken 71828 times.
|
97658 | if (cap > 0 && w == cap) |
| 484 | 3672 | break; | |
| 485 | } | ||
| 486 |
1/2✓ Branch 0 taken 276 times.
✗ Branch 1 not taken.
|
4246 | vs_idset_cleanup(&seen); |
| 487 | 3970 | return w; | |
| 488 | } | ||
| 489 | |||
| 490 | /* | ||
| 491 | * Partially order `results` so the `want` smallest-distance entries | ||
| 492 | * occupy the front (in arbitrary order within each side). Iterative | ||
| 493 | * Hoare quickselect with median-of-three pivots: expected O(n), no | ||
| 494 | * recursion, no external randomness. | ||
| 495 | */ | ||
| 496 | static void | ||
| 497 | 2819 | quickselect_by_distance(VsTopKEntry *results, uint32_t n, uint32_t want) | |
| 498 | { | ||
| 499 | 2819 | uint32_t lo = 0, hi = n; | |
| 500 | |||
| 501 |
6/6✓ Branch 0 taken 19003 times.
✓ Branch 1 taken 1498 times.
✓ Branch 2 taken 18929 times.
✓ Branch 3 taken 74 times.
✓ Branch 4 taken 17289 times.
✓ Branch 5 taken 1247 times.
|
20501 | while (hi - lo > 1 && want > lo && want < hi) |
| 502 | { | ||
| 503 | 17682 | uint32_t mid = lo + (hi - lo) / 2; | |
| 504 | 17682 | Distance a = results[lo].distance; | |
| 505 | 17682 | Distance b = results[mid].distance; | |
| 506 | 17682 | Distance c = results[hi - 1].distance; | |
| 507 | 393 | Distance pivot; | |
| 508 |
8/8✓ Branch 0 taken 9144 times.
✓ Branch 1 taken 8538 times.
✓ Branch 2 taken 5483 times.
✓ Branch 3 taken 3661 times.
✓ Branch 4 taken 8755 times.
✓ Branch 5 taken 4944 times.
✓ Branch 6 taken 3489 times.
✓ Branch 7 taken 5266 times.
|
17682 | if ((a <= b && b <= c) || (c <= b && b <= a)) |
| 509 | 7079 | pivot = b; | |
| 510 |
8/8✓ Branch 0 taken 5113 times.
✓ Branch 1 taken 5312 times.
✓ Branch 2 taken 2570 times.
✓ Branch 3 taken 2543 times.
✓ Branch 4 taken 5262 times.
✓ Branch 5 taken 2538 times.
✓ Branch 6 taken 2728 times.
✓ Branch 7 taken 2534 times.
|
10425 | else if ((b <= a && a <= c) || (c <= a && a <= b)) |
| 511 | 5138 | pivot = a; | |
| 512 | else | ||
| 513 | 5465 | pivot = c; | |
| 514 | |||
| 515 | /* The median-of-three pivot exists in the range, so each scan | ||
| 516 | * has a stopper; the explicit bounds checks make that local | ||
| 517 | * rather than an invariant to trust (j is unsigned, so an | ||
| 518 | * unchecked descent past lo would wrap). */ | ||
| 519 | 17682 | uint32_t i = lo, j = hi - 1; | |
| 520 |
2/2✓ Branch 0 taken 177288 times.
✓ Branch 1 taken 17682 times.
|
212652 | while (i <= j) |
| 521 | { | ||
| 522 |
4/4✓ Branch 0 taken 420352 times.
✓ Branch 1 taken 926 times.
✓ Branch 2 taken 243990 times.
✓ Branch 3 taken 176362 times.
|
421278 | while (i < hi - 1 && results[i].distance < pivot) |
| 523 | 243990 | i++; | |
| 524 |
4/4✓ Branch 0 taken 411462 times.
✓ Branch 1 taken 470 times.
✓ Branch 2 taken 234644 times.
✓ Branch 3 taken 176818 times.
|
411932 | while (j > lo && results[j].distance > pivot) |
| 525 | 234644 | j--; | |
| 526 |
2/2✓ Branch 0 taken 17419 times.
✓ Branch 1 taken 159869 times.
|
177288 | if (i <= j) |
| 527 | { | ||
| 528 | 165506 | VsTopKEntry tmp = results[i]; | |
| 529 | 165506 | results[i] = results[j]; | |
| 530 | 165506 | results[j] = tmp; | |
| 531 | 165506 | i++; | |
| 532 |
2/2✓ Branch 0 taken 5899 times.
✓ Branch 1 taken 159607 times.
|
165506 | if (j == 0) |
| 533 | ✗ | break; | |
| 534 | 165506 | j--; | |
| 535 | } | ||
| 536 | } | ||
| 537 | /* Entries [lo, j] <= pivot <= entries [i, hi). Recurse into the | ||
| 538 | * side containing the selection boundary. */ | ||
| 539 |
2/2✓ Branch 0 taken 9365 times.
✓ Branch 1 taken 8317 times.
|
17682 | if (want <= j + 1) |
| 540 | 9190 | hi = j + 1; | |
| 541 |
1/2✓ Branch 0 taken 8274 times.
✗ Branch 1 not taken.
|
8274 | else if (want >= i) |
| 542 | 8099 | lo = i; | |
| 543 | else | ||
| 544 | ✗ | return; /* boundary falls in the pivot-equal middle band */ | |
| 545 | } | ||
| 546 | } | ||
| 547 | |||
| 548 | /* ---------------------------------------------------------------- | ||
| 549 | * Extract sorted | ||
| 550 | * ---------------------------------------------------------------- */ | ||
| 551 | |||
| 552 | /* The uncapped form of vs_topk_extract_sorted_capped. */ | ||
| 553 | void | ||
| 554 | 3784 | vs_topk_extract_sorted(VsTopK *topk, VsTopKEntry *results, uint32_t *count_out) | |
| 555 | { | ||
| 556 | 3784 | vs_topk_extract_sorted_capped(topk, results, count_out, 0); | |
| 557 | 3784 | } | |
| 558 | |||
| 559 | /* | ||
| 560 | * Extract the threshold survivors sorted by distance ascending and | ||
| 561 | * deduplicated (first, i.e. best-distance, entry per id), keeping at | ||
| 562 | * most the best `cap` unique candidates; cap == 0 keeps them all. The | ||
| 563 | * survivor count is unbounded when distance estimates are noisy -- a | ||
| 564 | * loose threshold admits most scanned entries, every SOAR replica | ||
| 565 | * included -- so sorting all survivors is wasted work when only a | ||
| 566 | * small prefix is kept. A quickselect partition narrows to the | ||
| 567 | * smallest 3*cap entries first: every id occurs at most twice (primary | ||
| 568 | * posting plus at most one SOAR replica), so the smallest 2*cap | ||
| 569 | * entries already contain the best cap unique ids and 3*cap leaves | ||
| 570 | * margin. | ||
| 571 | */ | ||
| 572 | void | ||
| 573 | 7832 | vs_topk_extract_sorted_capped( | |
| 574 | VsTopK *topk, VsTopKEntry *results, uint32_t *count_out, uint32_t cap) | ||
| 575 | { | ||
| 576 |
2/2✓ Branch 0 taken 270 times.
✓ Branch 1 taken 12 times.
|
7832 | Distance threshold = vs_topk_threshold(topk); |
| 577 | |||
| 578 | 7832 | uint32_t out = 0; | |
| 579 |
2/2✓ Branch 0 taken 629061 times.
✓ Branch 1 taken 7832 times.
|
636893 | for (uint32_t i = 0; i < topk->cand_count; i++) |
| 580 | { | ||
| 581 | 629061 | Distance lb = topk->candidates[i].distance - topk->candidates[i].error; | |
| 582 |
2/2✓ Branch 0 taken 386403 times.
✓ Branch 1 taken 242658 times.
|
629061 | if (lb <= threshold) |
| 583 | 386403 | results[out++] = topk->candidates[i]; | |
| 584 | } | ||
| 585 | |||
| 586 | 7832 | uint32_t sel = out; | |
| 587 |
4/4✓ Branch 0 taken 4048 times.
✓ Branch 1 taken 3784 times.
✓ Branch 2 taken 2819 times.
✓ Branch 3 taken 1229 times.
|
7832 | if (cap > 0 && out > cap * 3) |
| 588 | { | ||
| 589 | 2819 | sel = cap * 3; | |
| 590 | 2819 | quickselect_by_distance(results, out, sel); | |
| 591 | } | ||
| 592 | |||
| 593 |
2/2✓ Branch 0 taken 4246 times.
✓ Branch 1 taken 3586 times.
|
7832 | if (sel > 1) |
| 594 | 4246 | qsort(results, sel, sizeof(VsTopKEntry), cmp_by_distance); | |
| 595 | |||
| 596 | 7832 | *count_out = dedup_sorted_prefix(results, sel, cap); | |
| 597 | 7832 | } | |
| 598 | |||
| 599 | /* | ||
| 600 | * Same as vs_topk_extract_sorted but skips the duplicate-id pass. | ||
| 601 | * For callers whose ids are unique by construction (array indices, | ||
| 602 | * beam-search buffer positions), the dedup scan can never remove | ||
| 603 | * anything and is quadratic in the survivor count. | ||
| 604 | */ | ||
| 605 | void | ||
| 606 | 902254 | vs_topk_extract_sorted_unique( | |
| 607 | VsTopK *topk, VsTopKEntry *results, uint32_t *count_out) | ||
| 608 | { | ||
| 609 |
2/2✓ Branch 0 taken 584460 times.
✓ Branch 1 taken 108272 times.
|
902254 | Distance threshold = vs_topk_threshold(topk); |
| 610 | |||
| 611 | /* Filter stale candidates and copy survivors to results */ | ||
| 612 | 902254 | uint32_t out = 0; | |
| 613 |
2/2✓ Branch 0 taken 27627413 times.
✓ Branch 1 taken 902254 times.
|
28529667 | for (uint32_t i = 0; i < topk->cand_count; i++) |
| 614 | { | ||
| 615 | 27627413 | Distance lb = topk->candidates[i].distance - topk->candidates[i].error; | |
| 616 |
2/2✓ Branch 0 taken 15481089 times.
✓ Branch 1 taken 12146324 times.
|
27627413 | if (lb <= threshold) |
| 617 | 15481089 | results[out++] = topk->candidates[i]; | |
| 618 | } | ||
| 619 | |||
| 620 | /* Sort by distance ascending */ | ||
| 621 |
2/2✓ Branch 0 taken 579024 times.
✓ Branch 1 taken 323230 times.
|
902254 | if (out > 1) |
| 622 | 579024 | qsort(results, out, sizeof(VsTopKEntry), cmp_by_distance); | |
| 623 | |||
| 624 | 902254 | *count_out = out; | |
| 625 | 902254 | } | |
| 626 |