| 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 | * Standalone: binary max-heap array for threshold tracking. | ||
| 60 | * PG: pairing heap via opaque pointer (defined in src/pg/topk.c). | ||
| 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 | } VsTopK; | ||
| 75 | |||
| 76 | /* ---------------------------------------------------------------- | ||
| 77 | * API | ||
| 78 | * ---------------------------------------------------------------- */ | ||
| 79 | |||
| 80 | /* | ||
| 81 | * Initialize a top-K collection. Creates a child memory context | ||
| 82 | * under the current context for internal buffers. Caller must be | ||
| 83 | * in a context with the desired lifetime. | ||
| 84 | * Use vs_topk_cleanup() to free. | ||
| 85 | */ | ||
| 86 | void vs_topk_init(VsTopK *topk, uint32_t k); | ||
| 87 | |||
| 88 | /* Free internal buffers (does not free the VsTopK struct itself). */ | ||
| 89 | void vs_topk_cleanup(VsTopK *topk); | ||
| 90 | |||
| 91 | /* | ||
| 92 | * Create a heap-allocated top-K collection. | ||
| 93 | * Returns NULL on allocation failure. | ||
| 94 | */ | ||
| 95 | VsTopK *vs_topk_create(uint32_t k); | ||
| 96 | |||
| 97 | /* Free a heap-allocated top-K collection. */ | ||
| 98 | void vs_topk_destroy(VsTopK *topk); | ||
| 99 | |||
| 100 | /* Reset to empty state (reuses existing storage). */ | ||
| 101 | void vs_topk_reset(VsTopK *topk); | ||
| 102 | |||
| 103 | /* | ||
| 104 | * Reset to empty state AND change k. Reuses the existing memory | ||
| 105 | * context (preserves the per-topk arena) but re-allocates ub_heap / | ||
| 106 | * ub_ids / candidates within it. Cheap compared to a full init + | ||
| 107 | * cleanup pair, because the memctx itself isn't created or destroyed. | ||
| 108 | * | ||
| 109 | * Use this when the same VsTopK is reused across calls that may | ||
| 110 | * want different k values (e.g. beam-search keeps beam_width | ||
| 111 | * candidates at intermediate levels, nprobe at the last). | ||
| 112 | */ | ||
| 113 | void vs_topk_reset_to_k(VsTopK *topk, uint32_t k); | ||
| 114 | |||
| 115 | /* | ||
| 116 | * Same as vs_topk_insert but skips the O(k) per-insert dedup scan. | ||
| 117 | * Use ONLY when the caller guarantees all ids are unique. The cluster | ||
| 118 | * scan (where SOAR / boundary replicas can collide) must keep using | ||
| 119 | * vs_topk_insert; centroid beam search and similar code that | ||
| 120 | * inserts each candidate exactly once should use this fast path. | ||
| 121 | */ | ||
| 122 | void vs_topk_insert_unique( | ||
| 123 | VsTopK *topk, Distance distance, Distance error, uint64_t id); | ||
| 124 | |||
| 125 | /* | ||
| 126 | * Insert a candidate. Pruned if lower_bound >= threshold. | ||
| 127 | * Otherwise updates the threshold heap and appends to the | ||
| 128 | * candidate buffer. | ||
| 129 | */ | ||
| 130 | void | ||
| 131 | vs_topk_insert(VsTopK *topk, Distance distance, Distance error, uint64_t id); | ||
| 132 | |||
| 133 | /* | ||
| 134 | * Current pruning threshold: the Kth-smallest upper bound | ||
| 135 | * (distance + error) seen so far. Returns INFINITY if fewer | ||
| 136 | * than K upper bounds have been recorded. | ||
| 137 | */ | ||
| 138 | static inline Distance | ||
| 139 | 56335087 | vs_topk_threshold(const VsTopK *topk) | |
| 140 | { | ||
| 141 |
10/10✓ Branch 0 taken 3522156 times.
✓ Branch 1 taken 10850033 times.
✓ Branch 2 taken 16772 times.
✓ Branch 3 taken 6667 times.
✓ Branch 4 taken 19230 times.
✓ Branch 5 taken 8326 times.
✓ Branch 6 taken 29223887 times.
✓ Branch 7 taken 12683617 times.
✓ Branch 8 taken 2998 times.
✓ Branch 9 taken 1401 times.
|
56335087 | if (topk->ub_count < topk->k) |
| 142 | 2929892 | return INFINITY; | |
| 143 | 40596367 | return topk->ub_heap[0]; | |
| 144 | } | ||
| 145 | |||
| 146 | /* | ||
| 147 | * Extract candidates sorted by distance ascending. Filters out | ||
| 148 | * stale entries whose lower bound now exceeds the final threshold. | ||
| 149 | * | ||
| 150 | * results[] must have space for topk->cand_count entries (upper | ||
| 151 | * bound; actual count returned via *count_out may be smaller). | ||
| 152 | * | ||
| 153 | * Does not reset the collection — call vs_topk_reset() or | ||
| 154 | * vs_topk_cleanup() when done with the results. | ||
| 155 | */ | ||
| 156 | void vs_topk_extract_sorted_capped( | ||
| 157 | VsTopK *topk, VsTopKEntry *results, uint32_t *count_out, uint32_t cap); | ||
| 158 | void vs_topk_extract_sorted( | ||
| 159 | VsTopK *topk, VsTopKEntry *results, uint32_t *count_out); | ||
| 160 | |||
| 161 | /* As above, but skips the duplicate-id pass; ids must be unique. */ | ||
| 162 | void vs_topk_extract_sorted_unique( | ||
| 163 | VsTopK *topk, VsTopKEntry *results, uint32_t *count_out); | ||
| 164 | |||
| 165 | #endif /* VS_TOPK_H */ | ||
| 166 |