| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | /* | ||
| 2 | * Copyright (c) 2026 Tiger Data, Inc. | ||
| 3 | * Licensed under the PostgreSQL License. See LICENSE for details. | ||
| 4 | * | ||
| 5 | * index_build.h - Shared index build utilities | ||
| 6 | * | ||
| 7 | * Generic helpers for building prism indexes, usable from both the | ||
| 8 | * PostgreSQL IAM build and the standalone CLI. All functions operate on the | ||
| 9 | * VsStorage abstraction and (where a tree is materialized at all) the | ||
| 10 | * HKMeansResult tree. | ||
| 11 | * | ||
| 12 | * The routing tree is the centroid tree: the hierarchy of centroid pages | ||
| 13 | * that routes queries and inserts down to the posting lists. These are the | ||
| 14 | * builders that write it straight to pages. | ||
| 15 | * | ||
| 16 | * Who runs what: the build has a serial shape and a parallel shape, and | ||
| 17 | * this header serves both. | ||
| 18 | * | ||
| 19 | * Serial (one process; do_serial_build in pg/build.c): | ||
| 20 | * prism_routing_tree_plan clusters the sample once (recording each | ||
| 21 | * node into a blob store) and sizes the page | ||
| 22 | * layout; | ||
| 23 | * prism_routing_tree_write replays the recorded nodes and streams | ||
| 24 | * every centroid + head page. | ||
| 25 | * | ||
| 26 | * Parallel (leader + N workers; see parallel_build.h for the barrier | ||
| 27 | * choreography). Clustering is divided by ROOT CHILD: after the | ||
| 28 | * cooperative sampling scan and the barrier-synchronized root k-means, | ||
| 29 | * the root's children are scheduled largest-first in batches of | ||
| 30 | * nparticipants, and each participant clusters one child's subtree per | ||
| 31 | * batch into its slot of a bounded DSM ring. Page writing is NOT | ||
| 32 | * divided: between the batch barriers the leader alone consumes each | ||
| 33 | * batch (recording layout counts, spilling the subtree blobs), and after | ||
| 34 | * the last batch the leader alone streams every page -- | ||
| 35 | * prism_routing_subtree_write writes one worker-built subtree's pages | ||
| 36 | * at its reserved block range and maps its | ||
| 37 | * local leaves into the global leaf index | ||
| 38 | * space (leaf_offset = prefix sum of the | ||
| 39 | * preceding children's leaf counts); | ||
| 40 | * prism_centroid_write_node writes the root page above the subtree | ||
| 41 | * roots. | ||
| 42 | * The posting scan that follows is cooperative again (workers route and | ||
| 43 | * encode into a shared sort; the leader merges), but that machinery | ||
| 44 | * lives in posting_build.h / parallel_build.h, not here. | ||
| 45 | * | ||
| 46 | * The standalone in-RAM build (no paging) uses prism_write_centroid_tree | ||
| 47 | * directly on a fully materialized tree. | ||
| 48 | */ | ||
| 49 | |||
| 50 | #ifndef PRISM_INDEX_BUILD_H | ||
| 51 | #define PRISM_INDEX_BUILD_H | ||
| 52 | |||
| 53 | struct PrismBlobStore; | ||
| 54 | #include <stdint.h> | ||
| 55 | |||
| 56 | #include "algo/hkmeans.h" | ||
| 57 | #include "core/types.h" | ||
| 58 | #include "index/centroid_build.h" | ||
| 59 | #include "index/centroid_page.h" | ||
| 60 | #include "index/centroid_search.h" /* PrismExactInternalCentroids */ | ||
| 61 | #include "index/storage.h" | ||
| 62 | #include "quant/rabitq.h" | ||
| 63 | |||
| 64 | /* ---------------------------------------------------------------- | ||
| 65 | * Exact internal-centroid collection (build-time routing accuracy) | ||
| 66 | * | ||
| 67 | * While the centroid tree streams to pages, the writer collects every | ||
| 68 | * INTERNAL node's float centroids — the exact values the pages' RaBitQ | ||
| 69 | * codes were encoded from — keyed by (page, entry). The build descent | ||
| 70 | * then scores the internal levels exactly through the | ||
| 71 | * PrismCentroidSearchState.exact_internal hook instead of the 1-bit | ||
| 72 | * estimates; the leaf level keeps the estimates and is exact-re-ranked | ||
| 73 | * per row via the head pages' full-precision pt_centroids. | ||
| 74 | * | ||
| 75 | * Addressing: internal and leaf-parent pages interleave in the centroid | ||
| 76 | * region (the serial write is post-order, the parallel write packs per | ||
| 77 | * subtree), so a single flat formula over the region would also carry | ||
| 78 | * the — vastly more numerous — leaf entries. Instead a per-page offset | ||
| 79 | * array maps each region page to the first slot of a dense | ||
| 80 | * internal-entry array (PRISM_EXACT_INTERNAL_NONE for leaf-parent pages). | ||
| 81 | * Within a node, entry e lands on page e / stride at page index | ||
| 82 | * e % stride, where stride is the format's fixed page capacity | ||
| 83 | * (prism_centroid_max_entries_fmt) — the exact packing every centroid | ||
| 84 | * writer uses, and each node starts on a fresh reserved page. | ||
| 85 | * | ||
| 86 | * Size: internal entries number nnodes - 1 (every non-root node is one | ||
| 87 | * parent entry), ~nlist / fan_out — a few thousand × dim floats even at | ||
| 88 | * very large nlist. | ||
| 89 | * ---------------------------------------------------------------- */ | ||
| 90 | typedef struct PrismExactCentroidCollector | ||
| 91 | { | ||
| 92 | Dimension dim; | ||
| 93 | BlockNumber base; /* first block of the centroid-page region */ | ||
| 94 | uint32_t npages; /* region length in pages */ | ||
| 95 | uint32_t stride; /* entries per full page for the format */ | ||
| 96 | uint32_t *page_off; /* [npages] first slot per page (or NONE) */ | ||
| 97 | float *cents; /* [cap * dim] collected centroids */ | ||
| 98 | uint32_t nslots; | ||
| 99 | uint32_t cap; | ||
| 100 | /* Collection budget: when the packed centroids would exceed it the | ||
| 101 | * collector latches overflowed, releases its arrays, and the | ||
| 102 | * collection/view degrade to empty — the build descent falls back to | ||
| 103 | * estimate-scored internal levels instead of allocating unbounded | ||
| 104 | * memory under adversarial nlist/fan_out configurations. */ | ||
| 105 | uint64_t max_bytes; | ||
| 106 | bool overflowed; | ||
| 107 | /* Dedicated context owning page_off and cents, created at init as | ||
| 108 | * a child of the caller's (build-scoped) current context. add_node | ||
| 109 | * grows cents during tree streaming, which runs inside a | ||
| 110 | * pass-scoped scratch context that dies before the collector's | ||
| 111 | * consumers (the build's encode scan and refine pass), so every | ||
| 112 | * collector allocation goes to this context — and cleanup is a | ||
| 113 | * single context delete, so nothing can dangle or double-free. */ | ||
| 114 | VsMemCtx ctx; | ||
| 115 | } PrismExactCentroidCollector; | ||
| 116 | |||
| 117 | /* | ||
| 118 | * Whether a build collects exact internal centroids: only multi-level | ||
| 119 | * trees have internal levels to score, and only the estimated centroid | ||
| 120 | * formats gain anything (FLOAT/HALF descents already score exactly). | ||
| 121 | * Shared by the serial and parallel builds so the two cannot drift. | ||
| 122 | */ | ||
| 123 | static inline bool | ||
| 124 | 277 | prism_exact_centroid_enabled(uint32_t nlevels, PrismCentroidFormat fmt) | |
| 125 | { | ||
| 126 |
6/6✓ Branch 0 taken 86 times.
✓ Branch 1 taken 191 times.
✓ Branch 2 taken 6 times.
✓ Branch 3 taken 24 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 4 times.
|
277 | return nlevels >= 2 && (fmt == PRISM_CENTROID_FMT_RABITQ || |
| 127 |
2/2✓ Branch 0 taken 32 times.
✓ Branch 1 taken 24 times.
|
56 | fmt == PRISM_CENTROID_FMT_FASTSCAN); |
| 128 | } | ||
| 129 | |||
| 130 | /* | ||
| 131 | * Collection budget from the build's memory budget (KB): an eighth — | ||
| 132 | * the collection is small (~nlist / fan_out centroids) next to the | ||
| 133 | * sort buffers, so the fraction only matters as a cap under | ||
| 134 | * adversarial nlist/fan_out. 0 (standalone default) means unbounded. | ||
| 135 | */ | ||
| 136 | static inline uint64_t | ||
| 137 | 258 | prism_exact_centroid_budget(uint64_t work_mem_kb) | |
| 138 | { | ||
| 139 |
5/6✓ Branch 0 taken 51 times.
✓ Branch 1 taken 26 times.
✓ Branch 2 taken 181 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 times.
✓ Branch 5 taken 180 times.
|
258 | return work_mem_kb > 0 ? work_mem_kb * 1024 / 8 : UINT64_MAX; |
| 140 | } | ||
| 141 | |||
| 142 | /* | ||
| 143 | * Estimated collection size for a planned tree shape: the collected | ||
| 144 | * slots are every non-root node, a geometric series over the levels | ||
| 145 | * that sums to roughly nlist / (fan_out - 1), off only by per-level | ||
| 146 | * rounding. Known from the build parameters alone, so an under-budget | ||
| 147 | * configuration can be reported at build start and the collector can | ||
| 148 | * pre-size its array; the collector still enforces the budget against | ||
| 149 | * the real size at collection time. | ||
| 150 | */ | ||
| 151 | static inline uint64_t | ||
| 152 | 260 | prism_exact_centroid_expected_slots(uint64_t nlist, uint32_t fan_out) | |
| 153 | { | ||
| 154 |
1/2✓ Branch 0 taken 79 times.
✗ Branch 1 not taken.
|
228 | return (fan_out > 1) ? nlist / (fan_out - 1) + 1 : nlist; |
| 155 | } | ||
| 156 | |||
| 157 | static inline uint64_t | ||
| 158 | 181 | prism_exact_centroid_expected_bytes( | |
| 159 | uint64_t nlist, uint32_t fan_out, Dimension dim) | ||
| 160 | { | ||
| 161 |
3/4✓ Branch 0 taken 149 times.
✓ Branch 1 taken 32 times.
✓ Branch 2 taken 181 times.
✗ Branch 3 not taken.
|
330 | return prism_exact_centroid_expected_slots(nlist, fan_out) * dim * |
| 162 | sizeof(float); | ||
| 163 | } | ||
| 164 | |||
| 165 | /* expected_slots pre-sizes the centroid array (0 = small default); | ||
| 166 | * pass prism_exact_centroid_expected_slots for the planned shape so | ||
| 167 | * growth is a rounding case, not the normal path. */ | ||
| 168 | void prism_exact_centroid_collector_init( | ||
| 169 | PrismExactCentroidCollector *c, | ||
| 170 | Dimension dim, | ||
| 171 | PrismCentroidFormat fmt, | ||
| 172 | BlockNumber base, | ||
| 173 | uint32_t npages, | ||
| 174 | uint64_t max_bytes, | ||
| 175 | uint64_t expected_slots); | ||
| 176 | |||
| 177 | /* Collect one internal node's centroids, written at first_blk. */ | ||
| 178 | void prism_exact_centroid_collector_add_node( | ||
| 179 | PrismExactCentroidCollector *c, | ||
| 180 | BlockNumber first_blk, | ||
| 181 | const float *cents, | ||
| 182 | uint32_t n); | ||
| 183 | |||
| 184 | void prism_exact_centroid_collector_cleanup(PrismExactCentroidCollector *c); | ||
| 185 | |||
| 186 | /* View straight over the collector's arrays (same-process use). */ | ||
| 187 | void prism_exact_centroid_view( | ||
| 188 | const PrismExactCentroidCollector *c, | ||
| 189 | PrismExactInternalCentroids *view); | ||
| 190 | |||
| 191 | /* | ||
| 192 | * Collection (de)serialization: the collector's sealed output, in the | ||
| 193 | * transport format described in index_build.c, for publishing to the | ||
| 194 | * parallel workers through the back-end seam (DSM segment / shared | ||
| 195 | * allocation). A NULL collector yields a valid empty collection whose | ||
| 196 | * view never matches any page (the hook stays inert) — used when | ||
| 197 | * collection is off, e.g. for exact centroid formats (FLOAT/HALF | ||
| 198 | * score exactly already). | ||
| 199 | */ | ||
| 200 | uint64_t | ||
| 201 | prism_exact_centroid_collection_size(const PrismExactCentroidCollector *c); | ||
| 202 | void prism_exact_centroid_collection_write( | ||
| 203 | const PrismExactCentroidCollector *c, void *collection); | ||
| 204 | void prism_exact_centroid_collection_view( | ||
| 205 | const void *collection, PrismExactInternalCentroids *view); | ||
| 206 | |||
| 207 | /* ---------------------------------------------------------------- | ||
| 208 | * Build statistics — shared between standalone and PG builds | ||
| 209 | * ---------------------------------------------------------------- */ | ||
| 210 | |||
| 211 | typedef struct PrismBuildStats | ||
| 212 | { | ||
| 213 | /* Phase timings (milliseconds) */ | ||
| 214 | double ms_total; /* total build time */ | ||
| 215 | double ms_sample; /* sampling vectors for clustering */ | ||
| 216 | double ms_kmeans; /* hierarchical k-means clustering */ | ||
| 217 | double ms_refine; /* full-table leaf-centroid refinement */ | ||
| 218 | double ms_setup; /* RaBitQ params, centroid rotation, etc */ | ||
| 219 | double ms_posting; /* posting build total (parallel + merge) */ | ||
| 220 | double ms_parallel; /* parallel encode+write phase */ | ||
| 221 | double ms_merge; /* serial partial page merge */ | ||
| 222 | double ms_centroid; /* writing centroid pages */ | ||
| 223 | |||
| 224 | /* Posting merge stats */ | ||
| 225 | uint32_t nworkers; | ||
| 226 | uint32_t total_pages; | ||
| 227 | uint32_t merge_input; | ||
| 228 | uint32_t merge_output; | ||
| 229 | } PrismBuildStats; | ||
| 230 | |||
| 231 | void prism_build_stats_print(const PrismBuildStats *s); | ||
| 232 | |||
| 233 | /* ---------------------------------------------------------------- | ||
| 234 | * Centroid layout + write | ||
| 235 | * ---------------------------------------------------------------- */ | ||
| 236 | |||
| 237 | /* | ||
| 238 | * Compute the block layout for centroid pages in a BFS tree. | ||
| 239 | * | ||
| 240 | * Assigns sequential block numbers to each BFS node's centroid | ||
| 241 | * pages, starting at first_blkno. Each node gets | ||
| 242 | * ceil(nchildren / max_entries) pages. | ||
| 243 | * | ||
| 244 | * node_first_blkno must have space for tree->nnodes entries. | ||
| 245 | * Returns the next available block number after all centroid pages. | ||
| 246 | */ | ||
| 247 | BlockNumber prism_compute_centroid_layout( | ||
| 248 | const HKMeansResult *tree, | ||
| 249 | uint32_t max_entries, | ||
| 250 | BlockNumber first_blkno, | ||
| 251 | BlockNumber *node_first_blkno); | ||
| 252 | |||
| 253 | /* | ||
| 254 | * Write centroid pages for all BFS nodes in the tree. | ||
| 255 | * | ||
| 256 | * Iterates the tree in BFS order, encoding each node's centroids | ||
| 257 | * and writing them to pages via the storage abstraction. Leaf | ||
| 258 | * nodes get PRISM_CENTROID_FLAG_LEAF; internal nodes get child_blkno | ||
| 259 | * pointers from node_first_blkno. Leaf entry j of a node gets the | ||
| 260 | * formula-derived posting head posting_base + node->first_leaf + j | ||
| 261 | * (no O(nlist) posting-head array). | ||
| 262 | * | ||
| 263 | * posting_base may be InvalidBlockNumber (centroid-only build without | ||
| 264 | * posting lists). | ||
| 265 | */ | ||
| 266 | /* | ||
| 267 | * pt_centroids: optional P^T * centroid array [nlist * dim] for leaf | ||
| 268 | * entries, stored alongside routing data. NULL to skip. | ||
| 269 | * collector: optional exact internal-centroid collection (NULL to skip). | ||
| 270 | */ | ||
| 271 | void prism_write_centroid_tree( | ||
| 272 | VsStorage *storage, | ||
| 273 | const HKMeansResult *tree, | ||
| 274 | Dimension dim, | ||
| 275 | uint32_t fan_out, | ||
| 276 | uint8_t level_offset, | ||
| 277 | PrismCentroidFormat centroid_format, | ||
| 278 | const RaBitQParams *rq_params, | ||
| 279 | const float *global_mean, | ||
| 280 | BlockNumber posting_base, | ||
| 281 | const BlockNumber *node_first_blkno, | ||
| 282 | const float *pt_centroids, | ||
| 283 | PrismExactCentroidCollector *collector); | ||
| 284 | |||
| 285 | /* ---------------------------------------------------------------- | ||
| 286 | * Streaming (page-backed) centroid-tree build | ||
| 287 | * | ||
| 288 | * Builds the hierarchical k-means tree top-down (DFS) and streams centroid | ||
| 289 | * pages straight to storage, never materializing the whole tree in RAM. Peak | ||
| 290 | * memory is the caller's sample buffer + an O(fan_out*depth*dim) recursion | ||
| 291 | * stack; the whole tree would be O(nlist*dim). Runs as two passes over the | ||
| 292 | * same recursion: a PLAN pass that clusters the sample once -- recording | ||
| 293 | * each node's clustering into the caller's blob store -- and reports the | ||
| 294 | * tree shape without writing; then a WRITE pass that replays the recorded | ||
| 295 | * nodes and emits the pages. Posting-list heads are formula-derived from | ||
| 296 | * the global leaf index, so the plan only sizes the page layout. | ||
| 297 | * | ||
| 298 | * These two entry points are the SERIAL build's whole clustering story | ||
| 299 | * (single process, no barriers). The parallel build divides the same work | ||
| 300 | * differently -- workers cluster per-root-child subtrees and only the | ||
| 301 | * leader writes pages, via prism_routing_subtree_write below -- but both | ||
| 302 | * shapes produce the identical page layout: reserved centroid blocks | ||
| 303 | * first, then the head region at first_posting, heads formula-derived. | ||
| 304 | * | ||
| 305 | * global_mean (the encoder centering) must be known before any page is | ||
| 306 | * written; the PLAN pass reports the mean of the leaf centroids | ||
| 307 | * (plan.leaf_mean) for that. The anchor is the leaf-centroid mean, not the | ||
| 308 | * per-vector sample mean: it centers the quantization on what the tree | ||
| 309 | * actually stores, and the quality of every centroid and posting code | ||
| 310 | * depends on it. | ||
| 311 | * ---------------------------------------------------------------- */ | ||
| 312 | |||
| 313 | /* | ||
| 314 | * Write one tree node's centroid page(s) in the node's format (fastscan or | ||
| 315 | * encoder-based). child_count is the non-leaf entries' child capacity; | ||
| 316 | * pass 0 for leaf-parent nodes. collector (optional, may be NULL) receives | ||
| 317 | * the node's exact float centroids when the node is internal (no | ||
| 318 | * PRISM_CENTROID_FLAG_LEAF in flags). | ||
| 319 | */ | ||
| 320 | void prism_centroid_write_node( | ||
| 321 | VsStorage *storage, | ||
| 322 | Dimension dim, | ||
| 323 | const float *cents, | ||
| 324 | uint32_t n, | ||
| 325 | PrismCentroidFormat fmt, | ||
| 326 | uint8_t level, | ||
| 327 | uint16_t flags, | ||
| 328 | uint16_t child_count, | ||
| 329 | const struct RaBitQParams *rq_params, | ||
| 330 | const float *global_mean, | ||
| 331 | const BlockNumber *child_blks, | ||
| 332 | const float *leaf_pt, | ||
| 333 | BlockNumber blkno, | ||
| 334 | PrismExactCentroidCollector *collector); | ||
| 335 | |||
| 336 | typedef struct PrismStreamTreePlan | ||
| 337 | { | ||
| 338 | uint32_t nleaves; | ||
| 339 | uint32_t nlevels; | ||
| 340 | uint32_t centroid_pages; /* pages the write pass will emit */ | ||
| 341 | float *leaf_mean; /* [dim] unweighted mean of the leaf centroids | ||
| 342 | * (vs_alloc; caller frees with vs_free) */ | ||
| 343 | } PrismStreamTreePlan; | ||
| 344 | |||
| 345 | /* | ||
| 346 | * PLAN pass: cluster the sample and report the tree shape (leaf count, | ||
| 347 | * depth, centroid page count, leaf-centroid mean) without writing anything. | ||
| 348 | * Returns false on k-means failure. | ||
| 349 | */ | ||
| 350 | /* | ||
| 351 | * Reserve the fixed page layout: extend the relation so blocks | ||
| 352 | * [0, end_blkno) exist before any is written. vs_storage_extend extends BY | ||
| 353 | * npages; the count doubles as the absolute layout end only because the | ||
| 354 | * relation holds nothing but the meta-page slot yet -- the reserved layout | ||
| 355 | * (heads at first_posting + leaf) silently shifts if a page ever sneaks in | ||
| 356 | * before this point, so the invariant is pinned here. | ||
| 357 | */ | ||
| 358 | static inline void | ||
| 359 | 276 | prism_build_reserve_layout(VsStorage *storage, BlockNumber end_blkno) | |
| 360 | { | ||
| 361 |
1/2✓ Branch 0 taken 194 times.
✗ Branch 1 not taken.
|
276 | BlockNumber ext_base = vs_storage_extend(storage, end_blkno); |
| 362 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 276 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
276 | Assert(ext_base == 0 || ext_base == InvalidBlockNumber); |
| 363 | 194 | (void)ext_base; | |
| 364 | 276 | } | |
| 365 | |||
| 366 | bool prism_routing_tree_plan( | ||
| 367 | const float *vectors, | ||
| 368 | uint32_t nvecs, | ||
| 369 | Dimension dim, | ||
| 370 | uint32_t nlist, | ||
| 371 | uint32_t fan_out, | ||
| 372 | DistanceMetric metric, | ||
| 373 | PrismCentroidFormat format, | ||
| 374 | const KMeansOptions *opts, | ||
| 375 | struct PrismBlobStore *store, | ||
| 376 | PrismStreamTreePlan *out); | ||
| 377 | |||
| 378 | /* | ||
| 379 | * Per-leaf callback fired during the write pass, once per leaf, with the | ||
| 380 | * leaf's global index and its float centroid (still resident at that moment). | ||
| 381 | * The build uses it to write each posting-list head page carrying pt_centroid | ||
| 382 | * = P^T*centroid — the exact encode reference — which cannot be recovered from | ||
| 383 | * a compressed centroid page afterward. May be NULL. | ||
| 384 | */ | ||
| 385 | typedef void (*PrismStreamLeafCb)( | ||
| 386 | void *arg, uint32_t leaf, const float *centroid); | ||
| 387 | |||
| 388 | /* | ||
| 389 | * WRITE pass: cluster the sample again (identical tree) and stream the | ||
| 390 | * centroid pages to `storage` via on-demand block allocation (post-order, root | ||
| 391 | * last). Leaf c's head is first_posting + c, and on_leaf (if set) fires per | ||
| 392 | * leaf so the caller can write that leaf's head page from the resident float | ||
| 393 | * centroid. Centroid pages occupy reserved blocks [first_centroid, | ||
| 394 | * first_centroid + plan.centroid_pages) post-order (root last), and posting | ||
| 395 | * heads live in the far posting area — so the caller must pre-extend the | ||
| 396 | * relation to cover both before calling. Returns the root block (the value the | ||
| 397 | * metadata page's first_centroid must carry), or InvalidBlockNumber on | ||
| 398 | * failure. | ||
| 399 | */ | ||
| 400 | BlockNumber prism_routing_tree_write( | ||
| 401 | VsStorage *storage, | ||
| 402 | uint32_t nvecs, | ||
| 403 | Dimension dim, | ||
| 404 | DistanceMetric metric, | ||
| 405 | uint32_t nlist, | ||
| 406 | uint32_t fan_out, | ||
| 407 | PrismCentroidFormat format, | ||
| 408 | const RaBitQParams *rq_params, | ||
| 409 | const float *global_mean, | ||
| 410 | struct PrismBlobStore *store, | ||
| 411 | BlockNumber first_posting, | ||
| 412 | BlockNumber first_centroid, | ||
| 413 | PrismStreamLeafCb on_leaf, | ||
| 414 | void *on_leaf_arg, | ||
| 415 | PrismExactCentroidCollector *collector); | ||
| 416 | |||
| 417 | /* | ||
| 418 | * Stream one already-built subtree (an HKMeansResult a parallel worker | ||
| 419 | * clustered into its ring slot, then spilled through the blob store) to | ||
| 420 | * centroid pages at its reserved block range, | ||
| 421 | * BFS layout (subtree root at first_block). Leader-only: workers cluster, | ||
| 422 | * the leader writes -- it calls this once per root child, replaying the | ||
| 423 | * blobs in the batch-schedule order while placing each at the child's own | ||
| 424 | * reserved range, so the whole tree is never materialized as one blob and | ||
| 425 | * page writing needs no cross-process coordination. | ||
| 426 | * | ||
| 427 | * The caller supplies the subtree's coordinates in the global layout, both | ||
| 428 | * prefix sums over the preceding children (from the PLAN pass's per-child | ||
| 429 | * counts): first_block for the block range, and leaf_offset for the leaf | ||
| 430 | * index space -- leaf local_leaf's head is first_posting + leaf_offset + | ||
| 431 | * local_leaf (formula-derived). level_offset places the subtree's | ||
| 432 | * relative node levels at their absolute tree depth (1 under the root). | ||
| 433 | * on_leaf fires per leaf with its float centroid so the caller can write | ||
| 434 | * the head page. Returns the subtree root block (== first_block). | ||
| 435 | */ | ||
| 436 | BlockNumber prism_routing_subtree_write( | ||
| 437 | VsStorage *storage, | ||
| 438 | const HKMeansResult *subtree, | ||
| 439 | Dimension dim, | ||
| 440 | DistanceMetric metric, | ||
| 441 | uint32_t fan_out, | ||
| 442 | uint8_t level_offset, | ||
| 443 | PrismCentroidFormat format, | ||
| 444 | const RaBitQParams *rq_params, | ||
| 445 | const float *global_mean, | ||
| 446 | BlockNumber first_posting, | ||
| 447 | uint32_t leaf_offset, | ||
| 448 | BlockNumber first_block, | ||
| 449 | PrismStreamLeafCb on_leaf, | ||
| 450 | void *on_leaf_arg, | ||
| 451 | PrismExactCentroidCollector *collector); | ||
| 452 | |||
| 453 | /* | ||
| 454 | * Auto-tune fan_out from nlist. | ||
| 455 | * | ||
| 456 | * When fan_out equals default_fan_out, derive a value that gives a | ||
| 457 | * balanced tree with ~nlist leaves. Uses sqrt for moderate nlist, | ||
| 458 | * cbrt for large nlist (> 65536). When fan_out has been | ||
| 459 | * explicitly set (differs from default), it is returned unchanged | ||
| 460 | * (capped to nlist if larger). | ||
| 461 | */ | ||
| 462 | uint32_t | ||
| 463 | prism_auto_fan_out(uint32_t fan_out, uint32_t nlist, uint32_t default_fan_out); | ||
| 464 | |||
| 465 | /* | ||
| 466 | * Default value of the target_pages index setting: posting pages a list | ||
| 467 | * should rest at. The target is in pages because a probe's cost is pages | ||
| 468 | * read -- a list thinner than a page still costs the page. | ||
| 469 | * | ||
| 470 | * Lives here, not with the other reloption bounds, so the value the option | ||
| 471 | * registers as its default and the value used for an index that sets no | ||
| 472 | * options cannot drift apart. | ||
| 473 | */ | ||
| 474 | #define PRISM_DEFAULT_TARGET_PAGES 3 | ||
| 475 | |||
| 476 | /* | ||
| 477 | * Floor on the per-list target, in entries: a centroid is only as good as | ||
| 478 | * the number of points it was fit from, which is independent of how those | ||
| 479 | * points pack into pages. Binds at high dimension, where a page holds few | ||
| 480 | * entries. | ||
| 481 | * | ||
| 482 | * Distinct from the build's k-means training samples per list (build.c, | ||
| 483 | * parallel_backend.c), which sizes the sample region; they share a value | ||
| 484 | * but must stay independently tunable. | ||
| 485 | */ | ||
| 486 | #define PRISM_MIN_ENTRIES_PER_LIST 256 | ||
| 487 | |||
| 488 | /* | ||
| 489 | * Target vectors per posting list at a given dimension: | ||
| 490 | * | ||
| 491 | * max(target_pages * entries_per_page(dim), PRISM_MIN_ENTRIES_PER_LIST) | ||
| 492 | * | ||
| 493 | * target_pages is the index's setting, or 0 to use its default. | ||
| 494 | * Entries per page is the continuation-page capacity; a list's first page | ||
| 495 | * holds fewer, carrying the float encode reference. | ||
| 496 | * | ||
| 497 | * The single source of truth for the intended list size: prism_auto_nlist | ||
| 498 | * derives nlist from it, prism_target_entries_per_list inverts that. | ||
| 499 | */ | ||
| 500 | uint32_t prism_target_entries_per_dim(Dimension dim, uint32_t target_pages); | ||
| 501 | |||
| 502 | /* | ||
| 503 | * Auto-tune nlist from a (possibly estimated) vector count: ~one list per | ||
| 504 | * prism_target_entries_per_dim(dim, target_pages) vectors, floored at | ||
| 505 | * sqrt(count) so small tables still get enough lists to build. Used by every | ||
| 506 | * build path when nlist is not set explicitly. The count source differs by | ||
| 507 | * back-end (reltuples / heap-block estimate in PostgreSQL, the in-memory | ||
| 508 | * vector count standalone). Back-ends clamp the result to their own nlist | ||
| 509 | * ceiling. | ||
| 510 | */ | ||
| 511 | uint32_t prism_auto_nlist(double count, Dimension dim, uint32_t target_pages); | ||
| 512 | |||
| 513 | /* | ||
| 514 | * Vectors per posting list at a given row count -- the resting size that | ||
| 515 | * maintenance should aim each list at. | ||
| 516 | * | ||
| 517 | * Above the dimension's target squared it is that target; below, the sqrt | ||
| 518 | * floor in prism_auto_nlist takes over and the size is ~sqrt(count), | ||
| 519 | * growing with the dataset. Hence a function, not a bare constant. | ||
| 520 | * | ||
| 521 | * Pass nlist 0 to derive the count from prism_auto_nlist. A non-zero nlist | ||
| 522 | * is used as given -- a direct caller can ask "entries per list if there | ||
| 523 | * were this many lists." PostgreSQL maintenance does not pass the reloption: | ||
| 524 | * honouring it would make a grown index never reach the split trigger, so | ||
| 525 | * resolve_target_entries() always passes 0. | ||
| 526 | */ | ||
| 527 | uint32_t prism_target_entries_per_list( | ||
| 528 | double count, uint32_t nlist, Dimension dim, uint32_t target_pages); | ||
| 529 | |||
| 530 | /* | ||
| 531 | * Find secondary cluster by plain distance (2nd-nearest centroid), | ||
| 532 | * given a distance-sorted candidate list (e.g. from a tree beam | ||
| 533 | * descent). The nearest candidate that is not primary_cluster is the | ||
| 534 | * 2nd-nearest centroid. | ||
| 535 | * | ||
| 536 | * Returns the secondary cluster index, or primary_cluster if the gap | ||
| 537 | * ratio exceeds epsilon (no replication needed). | ||
| 538 | * gap_ratio = (dist_2nd - dist_primary) / |dist_primary| | ||
| 539 | */ | ||
| 540 | uint32_t prism_find_secondary_cluster( | ||
| 541 | const uint32_t *cand_leaves, | ||
| 542 | const Distance *cand_dists, | ||
| 543 | uint32_t ncand, | ||
| 544 | uint32_t primary_cluster, | ||
| 545 | Distance primary_dist, | ||
| 546 | double epsilon); | ||
| 547 | |||
| 548 | /* | ||
| 549 | * Find secondary cluster via SOAR (Spilling with Orthogonality- | ||
| 550 | * Amplified Residuals). | ||
| 551 | * | ||
| 552 | * Computes the orthogonality-amplified distance for each searched centroid: | ||
| 553 | * OA(vec, c) = ||vec - c||^2 + lambda * dot(vec - c, r)^2 | ||
| 554 | * where r is the normalized residual from the primary centroid, and returns | ||
| 555 | * the centroid minimizing OA distance (excluding primary). When lambda=0 this | ||
| 556 | * degenerates to the standard 2nd-nearest. | ||
| 557 | * | ||
| 558 | * The search set is `count` leaves: the ids cand_leaves[0..count) when | ||
| 559 | * cand_leaves is non-NULL (the beam-descent candidates — O(count)), otherwise | ||
| 560 | * a full scan of leaves 0..count (pass count = nleaves). The candidate form is | ||
| 561 | * used in production because the SOAR optimum is always among the nearest | ||
| 562 | * leaves, so the result is unchanged while scaling to large nlist; NULL is for | ||
| 563 | * callers/tests that want the exhaustive scan. | ||
| 564 | */ | ||
| 565 | uint32_t prism_find_soar_secondary( | ||
| 566 | const float *vec, | ||
| 567 | const float *leaf_centroids, | ||
| 568 | const uint32_t *cand_leaves, | ||
| 569 | uint32_t count, | ||
| 570 | Dimension dim, | ||
| 571 | uint32_t primary_cluster, | ||
| 572 | const float *normalized_residual, | ||
| 573 | double lambda); | ||
| 574 | |||
| 575 | #endif /* PRISM_INDEX_BUILD_H */ | ||
| 576 |