GCC Code Coverage Report

Directory: src/
File: src/index/index_build.c
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 440 481 91.5%
Functions: 33 34 97.1%
Branches: 191 256 74.6%
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.c - Shared index build utilities
6 *
7 * Generic helpers for building prism indexes, usable from both
8 * the PostgreSQL IAM build and the standalone CLI.
9 */
10
11 #include <inttypes.h>
12 #include <math.h>
13 #include <string.h>
14
15 #include "algo/kmeans.h"
16 #include "algo/vecops.h"
17 #include "core/log.h"
18 #include "core/memory.h"
19 #include "index/centroid_build.h"
20 #include "index/index_build.h"
21 #include "index/parallel_build.h" /* PrismBlobStore seam (plan/write replay) */
22 #include "index/posting_page.h"
23 #include "quant/fastscan.h"
24
25 /* ----------------------------------------------------------------
26 * Exact internal-centroid collection — see index_build.h
27 * ---------------------------------------------------------------- */
28
29 /*
30 * A collection is the collector's sealed output: the exact centroids
31 * of the internal tree nodes, keyed by where each centroid's
32 * compressed twin sits on disk (centroid page, entry index on that
33 * page). Serialized as one contiguous allocation with no internal
34 * pointers, so it can be handed across the parallel-build seam (a DSM
35 * segment under PostgreSQL, a plain heap allocation standalone) and
36 * consumed in place by any backend that maps it.
37 *
38 * +--------------------------------+
39 * | ExactCentroidCollectionHeader | base, npages, nslots, dim
40 * +--------------------------------+
41 * | uint32_t page_off[npages] | slot base per page, or NONE
42 * +--------------------------------+
43 * | float cents[nslots * dim] | centroids, densely packed
44 * +--------------------------------+
45 *
46 * The collection covers the centroid-page region [base, base +
47 * npages): a
48 * page's page_off element is indexed by its offset from the starting
49 * page (blkno - base). Pages that carry internal entries hold the
50 * index of their first centroid in cents[]; their entries then map
51 * 1:1 and in page-entry order to the following slots, so entry i of
52 * page b scores against cents + (page_off[b - base] + i) * dim. Slot
53 * ranges of different pages are disjoint, but not ordered: cents[] is
54 * appended in tree-write order (post-order serial, per-subtree
55 * parallel), not block order. Leaf-level pages — collection skips
56 * them — hold PRISM_EXACT_INTERNAL_NONE instead.
57 *
58 * An empty collection (collecting off or over budget) is just the header
59 * with npages = 0 and base = InvalidBlockNumber, which no block can
60 * match — the consumer's scoring hook stays inert. nslots and dim are
61 * self-description (sizing and debugging); the view reads only base
62 * and npages, and slot extents come from page_off plus the pages' own
63 * entry counts.
64 *
65 * Lifecycle: the LEADER produces it — it alone streams the tree to
66 * pages, collecting into PrismExactCentroidCollector as it writes —
67 * then sizes and fills the collection from the collector and publishes it
68 * before the tree-ready barrier. The WORKERS attach after that
69 * barrier and wrap it in a PrismExactInternalCentroids view to score
70 * the internal tree levels while routing rows to their posting lists
71 * (and again in the refine pass); they only ever read it. The collection's
72 * memory is deliberately not context-managed and not the collector's:
73 * under PostgreSQL it is a DSM segment (process-shared, refcounted by
74 * the resource-owner machinery — a memory context is process-private
75 * and could not cross the seam), standalone a heap allocation the
76 * leader frees after the threads join. That split also lets the
77 * leader clean up the collector right after serializing, halving the
78 * peak, while the collection lives on until the last worker detaches.
79 * The serial build skips serialization entirely: same process, so its view
80 * points straight into the collector's arrays.
81 */
82 typedef struct ExactCentroidCollectionHeader
83 {
84 BlockNumber base;
85 uint32_t npages;
86 uint32_t nslots;
87 uint32_t dim;
88 } ExactCentroidCollectionHeader;
89
90 /* Fallback initial cents[] capacity (slots) when the caller passes no
91 * expected_slots pre-size; add_node doubles from here. Covers a
92 * typical two-level tree (fan_out + 1 nodes) without regrowing while
93 * staying well under any realistic budget. */
94 #define CENTROID_COLLECTOR_INIT_CAPACITY 256
95
96 void
97 83 prism_exact_centroid_collector_init(
98 PrismExactCentroidCollector *c,
99 Dimension dim,
100 PrismCentroidFormat fmt,
101 BlockNumber base,
102 uint32_t npages,
103 uint64_t max_bytes,
104 uint64_t expected_slots)
105 {
106 83 memset(c, 0, sizeof(*c));
107 83 c->dim = dim;
108 83 c->base = base;
109 83 c->npages = npages;
110 83 c->stride = prism_centroid_max_entries_fmt(dim, fmt);
111 /* Every allocation goes to the collector's dedicated child context
112 * (see the ownership contract on PrismExactCentroidCollector). */
113 83 c->ctx = vs_memctx_create(vs_memctx_current(), "vs exact centroids");
114 83 c->page_off = vs_memctx_alloc(c->ctx, (size_t)npages * sizeof(uint32_t));
115
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 51 times.
 83 memset(c->page_off, 0xFF, (size_t)npages * sizeof(uint32_t));
116 83 c->max_bytes = max_bytes;
117 /* Pre-size to the caller's estimate so growth is a rounding case,
118 * never past what the budget admits; an over-budget estimate is
119 * pointless to allocate in full -- collection will latch overflowed
120 * at the boundary anyway. */
121 134 uint64_t cap = expected_slots > 0 ? expected_slots
122
2/2
✓ Branch 0 taken 30 times.
✓ Branch 1 taken 53 times.
 83 : CENTROID_COLLECTOR_INIT_CAPACITY;
123 83 uint64_t max_cap = max_bytes / ((uint64_t)dim * sizeof(float));
124
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 30 times.
 83 if (cap > max_cap)
125 2 cap = max_cap;
126
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 32 times.
 83 if (cap == 0)
127 ✗ cap = 1;
128 83 c->cap = (uint32_t)cap;
129 83 c->cents = vs_memctx_alloc(c->ctx, (size_t)c->cap * dim * sizeof(float));
130 83 }
131
132 /*
133 * Ensure cents[] holds at least total slots. Reached only when the
134 * init pre-size undershot (per-level rounding in the estimate): grow
135 * geometrically — a fixed-step policy would copy O(n^2) bytes over
136 * the collection's lifetime, doubling copies less than 2x the final
137 * size in total — and clamp the capacity to the budget so the
138 * allocation never overshoots what the caller's budget check admits
139 * (that check guarantees the clamped capacity still fits).
140 */
141 static void
142 263 exact_centroid_collector_reserve(
143 PrismExactCentroidCollector *c, uint32_t total)
144 {
145
2/2
✓ Branch 0 taken 68 times.
✓ Branch 1 taken 195 times.
 263 if (total <= c->cap)
146 58 return;
147
148 20 uint64_t max_cap = c->max_bytes / ((uint64_t)c->dim * sizeof(float));
149
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
 20 Assert(total <= max_cap);
150 20 uint64_t new_cap = (uint64_t)c->cap * 2;
151
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 20 times.
 24 while (new_cap < total)
152 4 new_cap *= 2;
153
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
 20 if (new_cap > max_cap)
154 ✗ new_cap = max_cap;
155
156 /* Grow inside the collector's own context, not the streaming
157 * scratch context this is called under (that scratch dies at
158 * the end of the write pass, while the collector must survive
159 * until after the encode scan and the refine pass). The old
160 * array is freed eagerly (a no-op on the standalone arena) so
161 * both generations are never held at once. */
162 10 float *grown =
163 20 vs_memctx_alloc(c->ctx, (size_t)new_cap * c->dim * sizeof(float));
164 20 memcpy(grown, c->cents, (size_t)c->nslots * c->dim * sizeof(float));
165 20 vs_free(c->cents);
166 20 c->cents = grown;
167 20 c->cap = (uint32_t)new_cap;
168 }
169
170 void
171 276 prism_exact_centroid_collector_add_node(
172 PrismExactCentroidCollector *c,
173 BlockNumber first_blk,
174 const float *cents,
175 uint32_t n)
176 {
177
3/4
✓ Branch 0 taken 276 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 198 times.
✓ Branch 3 taken 78 times.
 276 if (n == 0 || c->overflowed)
178 2 return;
179
180 /* Enforce the collection budget: degrade to an empty (inert)
181 * collection rather than grow without bound. The arrays are dead
182 * once the flag latches (every consumer checks it first), so
183 * release them right away — the encode scan that follows is the
184 * build's memory peak. */
185 266 uint64_t need = ((uint64_t)c->nslots + n) * c->dim * sizeof(float);
186
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 263 times.
 266 if (need > c->max_bytes)
187 {
188
1/2
✓ Branch 1 taken 1 times.
✗ Branch 2 not taken.
 3 vs_warn("exact centroid collection over budget (needs more "
189 "than %" PRIu64 " of %" PRIu64 " bytes); build descent "
190 "falls back to estimated internal scoring — consider "
191 "raising the build memory budget",
192 need,
193 c->max_bytes);
194 3 c->overflowed = true;
195 3 prism_exact_centroid_collector_cleanup(c);
196 3 return;
197 }
198
199 263 exact_centroid_collector_reserve(c, c->nslots + n);
200
201 /* The node's entries pack stride per page from first_blk — the same
202 * packing every centroid writer uses (each node starts on a fresh
203 * reserved page). Record each page's slot base and append its
204 * centroids in page-entry order. */
205 263 uint32_t npages_node = (n + c->stride - 1) / c->stride;
206
2/2
✓ Branch 0 taken 263 times.
✓ Branch 1 taken 263 times.
 526 for (uint32_t p = 0; p < npages_node; p++)
207 {
208 263 BlockNumber blk = first_blk + p;
209
2/4
✓ Branch 0 taken 263 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 263 times.
 263 if (blk < c->base || blk - c->base >= c->npages)
210 {
211 /* Plan/write divergence: a node landed outside the planned
212 * centroid region. Degrade to inert rather than write
213 * page_off out of bounds in release builds. */
214 ✗ Assert(false);
215 c->overflowed = true;
216 prism_exact_centroid_collector_cleanup(c);
217 return;
218 }
219
220 263 uint32_t count = n - p * c->stride;
221
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 68 times.
 263 if (count > c->stride)
222 ✗ count = c->stride;
223
224 263 c->page_off[blk - c->base] = c->nslots;
225 263 memcpy(c->cents + (size_t)c->nslots * c->dim,
226 263 cents + (size_t)p * c->stride * c->dim,
227 263 (size_t)count * c->dim * sizeof(float));
228 263 c->nslots += count;
229 }
230 }
231
232 void
233 90 prism_exact_centroid_collector_cleanup(PrismExactCentroidCollector *c)
234 {
235 /* Everything the collector allocated lives in its dedicated
236 * context: one delete releases it all, exactly once (idempotent —
237 * a second call is a no-op). Callers must not run this until the
238 * last view consumer is done. */
239
2/2
✓ Branch 0 taken 83 times.
✓ Branch 1 taken 7 times.
 90 if (c->ctx != NULL)
240 83 vs_memctx_delete(c->ctx);
241 90 c->ctx = NULL;
242 90 c->page_off = NULL;
243 90 c->cents = NULL;
244 90 }
245
246 void
247 38 prism_exact_centroid_view(
248 const PrismExactCentroidCollector *c,
249 PrismExactInternalCentroids *view)
250 {
251
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 35 times.
 38 if (c->overflowed)
252 {
253 /* Inert view: a zero-page region matches no block. */
254 3 view->cents = NULL;
255 3 view->page_off = NULL;
256 3 view->base = InvalidBlockNumber;
257 3 view->npages = 0;
258 3 return;
259 }
260 35 view->cents = c->cents;
261 35 view->page_off = c->page_off;
262 35 view->base = c->base;
263 35 view->npages = c->npages;
264 }
265
266 uint64_t
267 133 prism_exact_centroid_collection_size(const PrismExactCentroidCollector *c)
268 {
269
4/4
✓ Branch 0 taken 51 times.
✓ Branch 1 taken 82 times.
✓ Branch 2 taken 21 times.
✓ Branch 3 taken 30 times.
 133 if (c == NULL || c->overflowed)
270 60 return sizeof(ExactCentroidCollectionHeader);
271 19 return sizeof(ExactCentroidCollectionHeader) +
272 49 (uint64_t)c->npages * sizeof(uint32_t) +
273 49 (uint64_t)c->nslots * c->dim * sizeof(float);
274 }
275
276 void
277 133 prism_exact_centroid_collection_write(
278 const PrismExactCentroidCollector *c, void *collection)
279 {
280 133 ExactCentroidCollectionHeader *hdr = (ExactCentroidCollectionHeader *)
281 collection;
282
283
4/4
✓ Branch 0 taken 51 times.
✓ Branch 1 taken 82 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 49 times.
 133 if (c == NULL || c->overflowed)
284 {
285 /* Empty collection: a zero-page region matches no block, so the view
286 * built from it leaves the scoring hook inert. */
287 84 hdr->base = InvalidBlockNumber;
288 84 hdr->npages = 0;
289 84 hdr->nslots = 0;
290 84 hdr->dim = 0;
291 84 return;
292 }
293
294 49 hdr->base = c->base;
295 49 hdr->npages = c->npages;
296 49 hdr->nslots = c->nslots;
297 49 hdr->dim = c->dim;
298
299 49 char *p = (char *)collection + sizeof(*hdr);
300 49 memcpy(p, c->page_off, (size_t)c->npages * sizeof(uint32_t));
301 49 p += (size_t)c->npages * sizeof(uint32_t);
302 49 memcpy(p, c->cents, (size_t)c->nslots * c->dim * sizeof(float));
303 }
304
305 void
306 198 prism_exact_centroid_collection_view(
307 const void *collection, PrismExactInternalCentroids *view)
308 {
309 198 const ExactCentroidCollectionHeader *hdr =
310 (const ExactCentroidCollectionHeader *)collection;
311 198 const char *p = (const char *)collection + sizeof(*hdr);
312
313 198 view->page_off = (const uint32_t *)p;
314 198 view->cents = (const float *)(p + (size_t)hdr->npages * sizeof(uint32_t));
315 198 view->base = hdr->base;
316 198 view->npages = hdr->npages;
317 198 }
318
319 /* Assign each node of a materialized tree its first block, packing nodes
320 * in index order from first_blkno; returns the block after the last. */
321 BlockNumber
322 640 prism_compute_centroid_layout(
323 const HKMeansResult *tree,
324 uint32_t max_entries,
325 BlockNumber first_blkno,
326 BlockNumber *node_first_blkno)
327 {
328 640 BlockNumber next = first_blkno;
329
330
2/2
✓ Branch 0 taken 1030 times.
✓ Branch 1 taken 640 times.
 1670 for (uint32_t i = 0; i < tree->nnodes; i++)
331 {
332
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 414 times.
 1030 uint32_t n = hk_nodes(tree)[i].nchildren;
333 1030 uint32_t npages = (n + max_entries - 1) / max_entries;
334
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1030 times.
 1030 if (npages == 0)
335 ✗ npages = 1;
336 1030 node_first_blkno[i] = next;
337 1030 next += npages;
338 }
339
340 640 return next;
341 }
342
343 /* Write a fully materialized (in-RAM) HKMeansResult tree to centroid
344 * pages. The streamed builders below replace this for the paged builds;
345 * the standalone in-RAM build still writes through it. */
346 void
347 327 prism_write_centroid_tree(
348 VsStorage *storage,
349 const HKMeansResult *tree,
350 Dimension dim,
351 uint32_t fan_out,
352 uint8_t level_offset,
353 PrismCentroidFormat centroid_format,
354 const RaBitQParams *rq_params,
355 const float *global_mean,
356 BlockNumber posting_base,
357 const BlockNumber *node_first_blkno,
358 const float *pt_centroids,
359 PrismExactCentroidCollector *collector)
360 {
361 /* Leaf child blocks are formula-derived (posting_base + global leaf
362 * index), so no O(nlist) posting-head array is needed. Each node has at
363 * most fan_out leaf entries, so this scratch is O(fan_out). */
364 462 BlockNumber *leaf_blks =
365 (posting_base != InvalidBlockNumber)
366 327 ? vs_alloc((size_t)fan_out * sizeof(BlockNumber))
367
1/2
✓ Branch 0 taken 327 times.
✗ Branch 1 not taken.
 327 : NULL;
368
369
2/2
✓ Branch 0 taken 547 times.
✓ Branch 1 taken 327 times.
 874 for (uint32_t i = 0; i < tree->nnodes; i++)
370 {
371
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 187 times.
 547 const HKMeansNode *node = &hk_nodes(tree)[i];
372 547 bool is_leaf = (node->level == tree->nlevels - 1);
373
374 547 uint16_t flags = is_leaf ? PRISM_CENTROID_FLAG_LEAF : 0;
375
2/2
✓ Branch 0 taken 314 times.
✓ Branch 1 taken 233 times.
 547 uint16_t child_count = is_leaf ? 0 : (uint16_t)fan_out;
376
377 207 const BlockNumber *child_blks;
378
3/4
✓ Branch 0 taken 481 times.
✓ Branch 1 taken 66 times.
✓ Branch 2 taken 294 times.
✗ Branch 3 not taken.
 547 if (is_leaf && leaf_blks != NULL)
379 {
380
2/2
✓ Branch 0 taken 2470 times.
✓ Branch 1 taken 481 times.
 2951 for (uint32_t j = 0; j < node->nchildren; j++)
381 2470 leaf_blks[j] = posting_base + node->first_leaf + j;
382 294 child_blks = leaf_blks;
383 }
384
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 else if (!is_leaf)
385 66 child_blks = &node_first_blkno[node->first_child];
386 else
387 ✗ child_blks = NULL;
388
389 /* Pass pt_centroids for leaf nodes only */
390
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 294 times.
 708 const float *leaf_pt = (is_leaf && pt_centroids != NULL)
391 ✗ ? pt_centroids +
392 ✗ (size_t)node->first_leaf * dim
393
2/2
✓ Branch 0 taken 294 times.
✓ Branch 1 taken 253 times.
 841 : NULL;
394
395 547 prism_centroid_write_node(
396 storage,
397 dim,
398 547 hk_node_centroids(tree, node),
399 547 node->nchildren,
400 centroid_format,
401 547 (uint8_t)(level_offset + node->level),
402 flags,
403 child_count,
404 rq_params,
405 global_mean,
406 child_blks,
407 leaf_pt,
408 547 node_first_blkno[i],
409 collector);
410 }
411
412
1/2
✓ Branch 0 taken 327 times.
✗ Branch 1 not taken.
 327 if (leaf_blks != NULL)
413 327 vs_free(leaf_blks);
414 327 }
415
416 /* ----------------------------------------------------------------
417 * Streaming (page-backed) centroid-tree build
418 * ---------------------------------------------------------------- */
419
420 typedef struct RoutingTreeCtx
421 {
422 /* inputs */
423 const float *vectors;
424 Dimension dim;
425 uint32_t nlist;
426 uint32_t fan_out;
427 uint32_t nlevels;
428 DistanceMetric metric;
429 KMeansOptions opts;
430 /* write-phase */
431 VsStorage *storage;
432 PrismCentroidFormat format;
433 const RaBitQParams *rq_params;
434 const float *global_mean;
435 BlockNumber first_posting; /* leaf c's head = first_posting + c */
436 PrismStreamLeafCb on_leaf;
437 void *on_leaf_arg;
438 /* write-phase: optional exact internal-centroid collection */
439 PrismExactCentroidCollector *collector;
440 /* plan-phase page-count helpers */
441 uint32_t max_ent; /* non-fastscan entries/page */
442 uint32_t fs_gpp; /* fastscan groups/page */
443 /* Node replay: the plan pass records each node's clustering
444 * ({k, centroids, assignments}, recursion order) here, and the write
445 * pass replays it instead of re-running k-means. Both passes require
446 * the store. */
447 PrismBlobStore *store;
448 float *replay_cents; /* write: [fan_out * dim] */
449 /* One node's uint16 assignments, shared by both passes (the plan pass
450 * narrows into it before recording, the replay pass reads records back
451 * into it). Sized once for the root -- the first and largest node --
452 * so it is never grown. */
453 uint16_t *assign_scratch; /* [nvecs] */
454 BlockNumber *leaf_blk_scratch; /* write: [fan_out] leaf head blocks */
455 /* accumulators */
456 uint32_t nleaves; /* running (== next first_leaf) */
457 double *leaf_sum; /* plan only, [dim]: leaf-centroid sum for the
458 * leaf_mean the caller uses as global_mean */
459 uint32_t centroid_pages; /* plan only */
460 BlockNumber next_blk; /* write only: next reserved centroid block */
461 bool ok;
462 } RoutingTreeCtx;
463
464 /*
465 * Write one tree node's centroid page(s) in the node's format: the fastscan
466 * writer carries its own encoder; every other format goes through the
467 * generic encoder + page writer. Shared by the streaming DFS and the
468 * parallel leader's root-page write.
469 */
470 void
471 2112 prism_centroid_write_node(
472 VsStorage *storage,
473 Dimension dim,
474 const float *cents,
475 uint32_t n,
476 PrismCentroidFormat fmt,
477 uint8_t level,
478 uint16_t flags,
479 uint16_t child_count,
480 const RaBitQParams *rq_params,
481 const float *global_mean,
482 const BlockNumber *child_blks,
483 const float *leaf_pt,
484 BlockNumber blkno,
485 PrismExactCentroidCollector *collector)
486 {
487 /* Every internal node of every build shape passes through here; the
488 * leaf level (LEAF flag) is exact-re-ranked per row instead. */
489
4/4
✓ Branch 0 taken 1549 times.
✓ Branch 1 taken 563 times.
✓ Branch 2 taken 270 times.
✓ Branch 3 taken 1279 times.
 2112 if (collector != NULL && (flags & PRISM_CENTROID_FLAG_LEAF) == 0)
490 270 prism_exact_centroid_collector_add_node(collector, blkno, cents, n);
491
492
2/2
✓ Branch 0 taken 1402 times.
✓ Branch 1 taken 710 times.
 2112 if (fmt == PRISM_CENTROID_FMT_FASTSCAN)
493 /* fastscan carries its own encoder; pt_centroids live on posting
494 * pages, never inline. */
495 1402 prism_centroid_write_fastscan_pages(
496 storage,
497 dim,
498 n,
499 level,
500 flags,
501 rq_params,
502 cents,
503 global_mean,
504 child_blks,
505 blkno);
506 else
507 {
508 152 CentroidEncoderState est;
509
3/4
✓ Branch 0 taken 50 times.
✓ Branch 1 taken 10 times.
✓ Branch 2 taken 92 times.
✗ Branch 3 not taken.
 710 CentroidEncoder *enc = centroid_encoder_init(
510 &est, fmt, cents, dim, rq_params, global_mean);
511 710 prism_centroid_write_pages(
512 storage,
513 dim,
514 n,
515 fmt,
516 level,
517 flags,
518 child_count,
519 enc,
520 child_blks,
521 leaf_pt,
522 blkno);
523 }
524 2112 }
525
526 /* Pages one node of `n` entries occupies — must match the writers' packing. */
527 static uint32_t
528 3032 node_npages(const RoutingTreeCtx *c, uint32_t n)
529 {
530
2/2
✓ Branch 0 taken 2304 times.
✓ Branch 1 taken 728 times.
 3032 if (c->format == PRISM_CENTROID_FMT_FASTSCAN)
531 {
532 2304 uint32_t ngroups = (n + VS_FASTSCAN_GROUP - 1) / VS_FASTSCAN_GROUP;
533
1/2
✓ Branch 0 taken 164 times.
✗ Branch 1 not taken.
 2304 uint32_t gpp = c->fs_gpp ? c->fs_gpp : 1;
534 2304 return (ngroups + gpp - 1) / gpp;
535 }
536
1/2
✓ Branch 0 taken 428 times.
✗ Branch 1 not taken.
 728 uint32_t me = c->max_ent ? c->max_ent : 1;
537 728 return (n + me - 1) / me;
538 }
539
540 /*
541 * Write one node's centroid page(s) at the next reserved block (post-order, so
542 * blocks are assigned in write order); returns the node's first block. The
543 * relation is pre-extended by the caller to cover the centroid + posting area,
544 * so writes use reserved blocks rather than appending (posting heads, which
545 * live in the far posting area, are written during this pass too).
546 */
547 static BlockNumber
548 1516 write_node_pages(
549 RoutingTreeCtx *c,
550 const float *cents,
551 uint32_t n,
552 uint32_t level,
553 bool is_leaf,
554 const BlockNumber *child_blks)
555 {
556 1516 uint16_t flags = is_leaf ? PRISM_CENTROID_FLAG_LEAF : 0;
557 1516 BlockNumber start = c->next_blk;
558
559 2845 prism_centroid_write_node(
560 c->storage,
561 1516 c->dim,
562 cents,
563 n,
564 c->format,
565 296 (uint8_t)level,
566 flags,
567 237 is_leaf ? 0 : (uint16_t)c->fan_out,
568
2/2
✓ Branch 0 taken 405 times.
✓ Branch 1 taken 1111 times.
 1516 c->rq_params,
569 c->global_mean,
570 child_blks,
571 NULL,
572 start,
573 c->collector);
574
575 1516 c->next_blk += node_npages(c, n);
576 1516 return start;
577 }
578
579 /*
580 * DFS one node. slice[count] are indices into c->vectors (NULL == identity for
581 * the root). Returns the node's first block (write phase) or
582 * InvalidBlockNumber (plan phase). Post-order: children are written before the
583 * parent so the parent's entries can carry their child block numbers.
584 */
585 /* Read exactly nbytes of a node record from the replay store. */
586 static bool
587 4548 replay_read(struct PrismBlobStore *store, void *dst, uint64_t nbytes)
588 {
589 4548 return prism_pbuild_blobstore_get(store, dst, nbytes) == nbytes;
590 }
591
592 /*
593 * Replay one node's clustering record from the blob store into the ctx's
594 * replay scratch. The plan pass wrote, per node in recursion order: the
595 * cluster count k (uint32), the centroids (k * dim floats) and the
596 * per-vector assignments (count uint16s); this reads the same three fields
597 * in the same order. Returns false on a short read or an out-of-range k,
598 * which can only mean the plan and write passes diverged.
599 */
600 static bool
601 1516 replay_node_record(RoutingTreeCtx *c, uint32_t count, uint32_t *out_k)
602 {
603 1516 uint32_t k = 0;
604
605
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 1516 times.
 1516 if (!replay_read(c->store, &k, sizeof(k)))
606 ✗ return false;
607
3/4
✓ Branch 0 taken 296 times.
✓ Branch 1 taken 1220 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1516 times.
 1516 if (k == 0 || k > c->fan_out)
608 ✗ return false;
609
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1516 times.
 1516 if (!replay_read(
610 1516 c->store,
611 1516 c->replay_cents,
612 1516 (uint64_t)k * c->dim * sizeof(float)))
613 ✗ return false;
614
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1516 times.
 1516 if (!replay_read(
615 1516 c->store,
616 1516 c->assign_scratch,
617 1516 (uint64_t)count * sizeof(uint16_t)))
618 ✗ return false;
619
620 1516 *out_k = k;
621 1516 return true;
622 }
623
624 /*
625 * Record one node's clustering into the blob store, in the exact field
626 * order replay_node_record reads back: k, centroids[k * dim],
627 * assignments[count]. Assignments are narrowed to uint16 -- a node
628 * clusters into at most fan_out (<= the meta page's 8-bit fan_out) groups.
629 */
630 static void
631 1516 record_node_clustering(
632 RoutingTreeCtx *c, const KMeansResult *km, uint32_t count)
633 {
634 1516 const size_t assign_nbytes = (size_t)count * sizeof(uint16_t);
635 1516 uint16_t *a16 = c->assign_scratch;
636
637
2/2
✓ Branch 0 taken 528513 times.
✓ Branch 1 taken 1516 times.
 530029 for (uint32_t v = 0; v < count; v++)
638 528513 a16[v] = (uint16_t)km->assignments[v];
639
640 1516 prism_pbuild_blobstore_put(c->store, &km->nlist, sizeof(km->nlist));
641 1516 prism_pbuild_blobstore_put(
642 c->store,
643 1516 km->centroids,
644 1516 (uint64_t)km->nlist * c->dim * sizeof(float));
645 1516 prism_pbuild_blobstore_put(c->store, a16, assign_nbytes);
646 1516 }
647
648 /* Compact the non-empty clusters' centroids to the front, preserving
649 * ascending cluster order (descent indexes the kept centroids densely). */
650 static uint32_t
651 1279 compact_nonempty_centroids(
652 float *cents,
653 const uint32_t *counts,
654 uint32_t nclusters,
655 Dimension dim)
656 {
657 1279 uint32_t kept = 0;
658
2/2
✓ Branch 0 taken 11093 times.
✓ Branch 1 taken 1279 times.
 12372 for (uint32_t cl = 0; cl < nclusters; cl++)
659 {
660
2/2
✓ Branch 0 taken 88 times.
✓ Branch 1 taken 11005 times.
 11093 if (counts[cl] == 0)
661 88 continue;
662
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 10981 times.
 11005 if (kept != cl)
663 24 memcpy(cents + (size_t)kept * dim,
664 24 cents + (size_t)cl * dim,
665 24 (size_t)dim * sizeof(float));
666 11005 kept++;
667 }
668 1279 return kept;
669 }
670
671 static void plan_node_recurse(
672 RoutingTreeCtx *c,
673 const uint32_t *slice,
674 uint32_t count,
675 uint32_t level);
676 static BlockNumber
677 replay_node_recurse(RoutingTreeCtx *c, uint32_t count, uint32_t level);
678
679 /* Leaf parent, plan side: the non-empty clusters become leaves. Nothing is
680 * written; count the leaf pages, tally the leaves, and fold the kept
681 * centroids into the leaf sum (the global mean's numerator). */
682 static void
683 1279 plan_leaf_parent(
684 RoutingTreeCtx *c, const KMeansResult *km, const uint32_t *counts)
685 {
686 1279 uint32_t kept = 0;
687
2/2
✓ Branch 0 taken 11093 times.
✓ Branch 1 taken 1279 times.
 12372 for (uint32_t cl = 0; cl < km->nlist; cl++)
688 {
689
2/2
✓ Branch 0 taken 88 times.
✓ Branch 1 taken 11005 times.
 11093 if (counts[cl] == 0)
690 88 continue;
691
1/2
✓ Branch 0 taken 11005 times.
✗ Branch 1 not taken.
 11005 if (c->leaf_sum != NULL)
692 {
693 11005 const float *lc = km->centroids + (size_t)cl * c->dim;
694
2/2
✓ Branch 0 taken 755215 times.
✓ Branch 1 taken 11005 times.
 766220 for (Dimension d = 0; d < c->dim; d++)
695 755215 c->leaf_sum[d] += lc[d];
696 }
697 11005 kept++;
698 }
699 1279 c->centroid_pages += node_npages(c, kept);
700 1279 c->nleaves += kept;
701 1279 }
702
703 /* Internal node, plan side: gather each non-empty child's slice of the
704 * sample and recurse, in ascending cluster order -- the recursion order
705 * that frames the tape -- then count this node's page(s). */
706 static void
707 237 plan_internal_node(
708 RoutingTreeCtx *c,
709 const KMeansResult *km,
710 const uint32_t *counts,
711 const uint32_t *slice,
712 uint32_t count,
713 uint32_t level)
714 {
715 237 uint32_t kept = 0;
716
3/4
✓ Branch 0 taken 1344 times.
✓ Branch 1 taken 237 times.
✓ Branch 2 taken 1344 times.
✗ Branch 3 not taken.
 1581 for (uint32_t cl = 0; cl < km->nlist && c->ok; cl++)
717 {
718
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 1341 times.
 1344 if (counts[cl] == 0)
719 3 continue;
720
721 1341 uint32_t sub_n = counts[cl];
722 1341 uint32_t *sub = vs_alloc((size_t)sub_n * sizeof(uint32_t));
723 1341 uint32_t idx = 0;
724
3/3
✓ Branch 0 taken 76800 times.
✓ Branch 1 taken 2313510 times.
✓ Branch 2 taken 1069 times.
 2391379 for (uint32_t v = 0; v < count; v++)
725
2/2
✓ Branch 0 taken 268173 times.
✓ Branch 1 taken 2121865 times.
 2390038 if (km->assignments[v] == cl)
726
2/2
✓ Branch 0 taken 59723 times.
✓ Branch 1 taken 208450 times.
 268173 sub[idx++] = slice ? slice[v] : v;
727 1341 plan_node_recurse(c, sub, sub_n, level + 1);
728 1341 vs_free(sub);
729 1341 kept++;
730 }
731 237 c->centroid_pages += node_npages(c, kept);
732 237 }
733
734 /*
735 * Plan pass -- writes the tape. Recursive DFS from the given node down
736 * (the mutual recursion runs through plan_internal_node): cluster the
737 * node's slice of the sample (slice[count] indexes c->vectors; NULL =
738 * identity at the root), record the clustering for the replay pass,
739 * descend into the non-empty children, and accumulate the layout totals
740 * (centroid pages, leaves, the leaf-centroid sum) that must be known
741 * before any page can be written.
742 */
743 static void
744 1516 plan_node_recurse(
745 RoutingTreeCtx *c,
746 const uint32_t *slice,
747 uint32_t count,
748 uint32_t level)
749 {
750
2/2
✓ Branch 0 taken 1220 times.
✓ Branch 1 taken 296 times.
 1516 if (!c->ok)
751 ✗ return;
752
753
2/2
✓ Branch 0 taken 123 times.
✓ Branch 1 taken 1393 times.
 1516 uint32_t k = (c->nlevels == 1) ? c->nlist : c->fan_out;
754
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 296 times.
 1516 if (k > count)
755 ✗ k = count;
756
757 2736 KMeansResult *km = vs_kmeans(
758 1516 c->vectors,
759 slice,
760 VS_VEC_F32,
761 count,
762 1516 c->dim,
763 k,
764 c->metric,
765 1516 &c->opts);
766
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1516 times.
 1516 if (km == NULL)
767 {
768 ✗ c->ok = false;
769 ✗ return;
770 }
771 1516 c->opts.initial_centroids = NULL; /* root only (matches vs_hkmeans_f32) */
772
773
1/2
✓ Branch 0 taken 1516 times.
✗ Branch 1 not taken.
 1516 if (c->store != NULL)
774 1516 record_node_clustering(c, km, count);
775
776 /* Count assignments per cluster (km->cluster_sizes may be stale). */
777 1516 uint32_t *counts = vs_alloc0((size_t)km->nlist * sizeof(uint32_t));
778
3/3
✓ Branch 0 taken 31200 times.
✓ Branch 1 taken 497609 times.
✓ Branch 2 taken 1220 times.
 530029 for (uint32_t v = 0; v < count; v++)
779 528513 counts[km->assignments[v]]++;
780
781
2/2
✓ Branch 0 taken 1279 times.
✓ Branch 1 taken 237 times.
 1516 if (level == c->nlevels - 1)
782 1279 plan_leaf_parent(c, km, counts);
783 else
784 237 plan_internal_node(c, km, counts, slice, count, level);
785
786 1516 vs_free(counts);
787 1516 vs_kmeans_result_destroy(km);
788 }
789
790 /* Leaf parent, replay side: compact the non-empty clusters into leaves,
791 * write the node's page(s) with the formula-derived leaf head blocks, and
792 * emit each leaf's head page from its resident float centroid. */
793 static BlockNumber
794 1279 replay_leaf_parent(
795 RoutingTreeCtx *c,
796 uint32_t nclusters,
797 const uint32_t *counts,
798 uint32_t level)
799 {
800 1279 float *cents = c->replay_cents;
801 1047 uint32_t kept =
802 1279 compact_nonempty_centroids(cents, counts, nclusters, c->dim);
803
804 /* Cosine routes and encodes against unit-norm references; k-means
805 * means drift below unit norm, so normalize the leaf centroids in
806 * place before the leaf codes and head references are written from
807 * them (the flat-tree path already does). */
808
2/2
✓ Branch 0 taken 25 times.
✓ Branch 1 taken 1254 times.
 1279 if (c->metric == DISTANCE_COSINE)
809
2/2
✓ Branch 0 taken 238 times.
✓ Branch 1 taken 25 times.
 263 for (uint32_t kk = 0; kk < kept; kk++)
810 238 vs_l2_normalize(cents + (size_t)kk * c->dim, c->dim);
811
812 /* Leaf heads are formula-derived (first_posting + global leaf index);
813 * kept <= fan_out, so the pass-lifetime scratch covers every node. */
814 1279 BlockNumber *leaf_blks = NULL;
815
1/2
✓ Branch 0 taken 1279 times.
✗ Branch 1 not taken.
 1279 if (c->first_posting != InvalidBlockNumber)
816 {
817 1279 leaf_blks = c->leaf_blk_scratch;
818
2/2
✓ Branch 0 taken 11005 times.
✓ Branch 1 taken 1279 times.
 12284 for (uint32_t kk = 0; kk < kept; kk++)
819 11005 leaf_blks[kk] = c->first_posting + c->nleaves + kk;
820 }
821 1279 BlockNumber blk = write_node_pages(c, cents, kept, level, true, leaf_blks);
822
823
1/2
✓ Branch 0 taken 1279 times.
✗ Branch 1 not taken.
 1279 if (c->on_leaf != NULL)
824
2/2
✓ Branch 0 taken 11005 times.
✓ Branch 1 taken 1279 times.
 12284 for (uint32_t kk = 0; kk < kept; kk++)
825 11005 c->on_leaf(
826 c->on_leaf_arg,
827 11005 c->nleaves + kk,
828 11005 cents + (size_t)kk * c->dim);
829 1279 c->nleaves += kept;
830 1279 return blk;
831 }
832
833 /* Internal node, replay side: recurse into each non-empty child in
834 * ascending cluster order -- the same order the plan pass recursed, which
835 * is what keeps the tape in frame -- then write this node's page(s) above
836 * the children. The children's records overwrite the shared replay
837 * scratch, so this node's centroids are snapshotted first. */
838 static BlockNumber
839 237 replay_internal_node(
840 RoutingTreeCtx *c,
841 uint32_t nclusters,
842 const uint32_t *counts,
843 uint32_t level)
844 {
845 237 const size_t vec_nbytes = (size_t)c->dim * sizeof(float);
846 237 const size_t cents_nbytes = (size_t)nclusters * vec_nbytes;
847
848 237 float *my_cents = vs_alloc(cents_nbytes);
849 237 memcpy(my_cents, c->replay_cents, cents_nbytes);
850
851 237 BlockNumber *child_blocks = vs_alloc(
852 64 (size_t)nclusters * sizeof(BlockNumber));
853 237 uint32_t kept = 0;
854
4/5
✓ Branch 0 taken 272 times.
✓ Branch 1 taken 1136 times.
✓ Branch 2 taken 445 times.
✓ Branch 3 taken 1072 times.
✗ Branch 4 not taken.
 1581 for (uint32_t cl = 0; cl < nclusters && c->ok; cl++)
855 {
856
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 1341 times.
 1344 if (counts[cl] == 0)
857 3 continue;
858
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1341 times.
 1341 if (kept != cl)
859 ✗ memcpy(my_cents + (size_t)kept * c->dim,
860 ✗ my_cents + (size_t)cl * c->dim,
861 vec_nbytes);
862 1341 child_blocks[kept] = replay_node_recurse(c, counts[cl], level + 1);
863 1341 kept++;
864 }
865
866 237 BlockNumber blk = InvalidBlockNumber;
867
1/2
✓ Branch 0 taken 237 times.
✗ Branch 1 not taken.
 237 if (c->ok)
868 237 blk = write_node_pages(c, my_cents, kept, level, false, child_blocks);
869 237 vs_free(child_blocks);
870 237 vs_free(my_cents);
871 237 return blk;
872 }
873
874 /*
875 * Replay pass -- reads the tape. Recursive DFS from the given node down
876 * (the mutual recursion runs through replay_internal_node): consume the
877 * record the plan pass wrote for this node and write its page(s). The
878 * record carries the centroids and the assignments, and each child's size
879 * falls out of the assignment counts, so unlike the plan walker this one
880 * needs no sample vectors and descends on counts alone. Post-order:
881 * children are written before the parent so the parent's entries can carry
882 * their block numbers. Returns the node's first block.
883 */
884 static BlockNumber
885 1516 replay_node_recurse(RoutingTreeCtx *c, uint32_t count, uint32_t level)
886 {
887
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1516 times.
 1516 if (!c->ok)
888 ✗ return InvalidBlockNumber;
889
890 1516 uint32_t nclusters = 0;
891
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 1516 times.
 1516 if (!replay_node_record(c, count, &nclusters))
892 {
893 ✗ c->ok = false;
894 ✗ return InvalidBlockNumber;
895 }
896
897 /* Count assignments per cluster. */
898 1516 uint32_t *counts = vs_alloc0((size_t)nclusters * sizeof(uint32_t));
899
3/3
✓ Branch 0 taken 31200 times.
✓ Branch 1 taken 497609 times.
✓ Branch 2 taken 1220 times.
 530029 for (uint32_t v = 0; v < count; v++)
900 528513 counts[c->assign_scratch[v]]++;
901
902 1220 BlockNumber blk;
903
2/2
✓ Branch 0 taken 1279 times.
✓ Branch 1 taken 237 times.
 1516 if (level == c->nlevels - 1)
904 1279 blk = replay_leaf_parent(c, nclusters, counts, level);
905 else
906 237 blk = replay_internal_node(c, nclusters, counts, level);
907
908 1516 vs_free(counts);
909 1516 return blk;
910 }
911
912 /* Shared walker state: tree geometry and page capacities. The plan entry
913 * adds the sample/metric/k-means knobs; the replay entry adds the storage,
914 * tape and replay scratch. */
915 static void
916 350 routing_tree_ctx_init(
917 RoutingTreeCtx *c,
918 Dimension dim,
919 uint32_t nlist,
920 uint32_t fan_out,
921 PrismCentroidFormat format)
922 {
923 350 memset(c, 0, sizeof(*c));
924 350 c->dim = dim;
925 350 c->nlist = nlist;
926 350 c->fan_out = fan_out < 2 ? 2 : fan_out;
927 350 c->nlevels = vs_hkmeans_nlevels(nlist, c->fan_out);
928 350 c->format = format;
929 350 c->max_ent = prism_centroid_max_entries_fmt(dim, format);
930 350 c->fs_gpp = prism_centroid_fastscan_max_groups(dim);
931 350 c->ok = true;
932 350 }
933
934 /* Plan-pass entry ("write the tape"): cluster the sample once, record
935 * every node, and return the layout totals -- see the header comment. */
936 bool
937 175 prism_routing_tree_plan(
938 const float *vectors,
939 uint32_t nvecs,
940 Dimension dim,
941 uint32_t nlist,
942 uint32_t fan_out,
943 DistanceMetric metric,
944 PrismCentroidFormat format,
945 const KMeansOptions *opts,
946 PrismBlobStore *store,
947 PrismStreamTreePlan *out)
948 {
949 151 RoutingTreeCtx c;
950 175 routing_tree_ctx_init(&c, dim, nlist, fan_out, format);
951 /* Only the plan pass clusters, so only it carries the sample, the
952 * metric and the k-means options. */
953 175 c.vectors = vectors;
954 175 c.metric = metric;
955
1/2
✓ Branch 0 taken 175 times.
✗ Branch 1 not taken.
 175 c.opts = opts ? *opts : (KMeansOptions)VS_KMEANS_OPTIONS_DEFAULT;
956 175 c.store = store;
957 175 c.leaf_sum = vs_alloc0((size_t)dim * sizeof(double));
958
1/2
✓ Branch 0 taken 175 times.
✗ Branch 1 not taken.
 175 if (store != NULL)
959 {
960 /* Recording narrows each node's assignments into this scratch; the
961 * root -- first and largest -- sizes it for the whole pass. */
962 175 c.assign_scratch = vs_alloc((size_t)nvecs * sizeof(uint16_t));
963 }
964
965 /*
966 * The recursion's per-node scratch (k-means temporaries, index slices,
967 * child-block arrays) lives in its own context so anything a node fails
968 * to free is reclaimed here, at the end of the pass, rather than
969 * accumulating for the rest of the build. This pass-scoped context is
970 * the deliberate leak-containment boundary: node lifetimes nest (a
971 * parent's k-means result and child arrays stay live across its
972 * children), so finer-grained reclamation such as resets at node or
973 * sibling boundaries would free live ancestor state — and would turn a
974 * missed free (a bounded, observable leak) into a use-after-free.
975 * Within a node, explicit frees remain the mechanism; the context caps
976 * their blast radius at one pass. The output (leaf_mean) is allocated
977 * in the caller's context outside the switch.
978 */
979 175 VsMemCtx scratch = vs_memctx_create(NULL, "vs stream plan");
980 175 VsMemCtx old_ctx = vs_memctx_switch(scratch);
981 175 plan_node_recurse(&c, NULL, nvecs, 0);
982 175 vs_memctx_switch(old_ctx);
983 175 vs_memctx_delete(scratch);
984
985
1/2
✓ Branch 0 taken 175 times.
✗ Branch 1 not taken.
 175 if (c.assign_scratch != NULL)
986 175 vs_free(c.assign_scratch);
987
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 175 times.
 175 if (!c.ok)
988 {
989 ✗ vs_free(c.leaf_sum);
990 ✗ return false;
991 }
992
993 175 out->nleaves = c.nleaves;
994 175 out->nlevels = c.nlevels;
995 175 out->centroid_pages = c.centroid_pages;
996 175 out->leaf_mean = vs_alloc((size_t)dim * sizeof(float));
997
2/2
✓ Branch 0 taken 15703 times.
✓ Branch 1 taken 175 times.
 15878 for (Dimension d = 0; d < dim; d++)
998 15703 out->leaf_mean[d] = c.nleaves > 0
999 15703 ? (float)(c.leaf_sum[d] / (double)c.nleaves)
1000
1/2
✓ Branch 0 taken 15703 times.
✗ Branch 1 not taken.
 15703 : 0.0f;
1001 175 vs_free(c.leaf_sum);
1002 175 return true;
1003 }
1004
1005 /* Replay-pass entry ("read the tape"): stream every centroid + head page
1006 * from the recorded nodes; returns the root block. */
1007 BlockNumber
1008 175 prism_routing_tree_write(
1009 VsStorage *storage,
1010 uint32_t nvecs,
1011 Dimension dim,
1012 DistanceMetric metric,
1013 uint32_t nlist,
1014 uint32_t fan_out,
1015 PrismCentroidFormat format,
1016 const RaBitQParams *rq_params,
1017 const float *global_mean,
1018 PrismBlobStore *store,
1019 BlockNumber first_posting,
1020 BlockNumber first_centroid,
1021 PrismStreamLeafCb on_leaf,
1022 void *on_leaf_arg,
1023 PrismExactCentroidCollector *collector)
1024 {
1025 /* Replay never clusters: everything it needs is on the tape, so it
1026 * takes no sample vectors, metric or k-means options. */
1027
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 175 times.
 175 if (store == NULL)
1028 ✗ return InvalidBlockNumber;
1029
1030 151 RoutingTreeCtx c;
1031 175 routing_tree_ctx_init(&c, dim, nlist, fan_out, format);
1032 175 c.metric = metric;
1033 175 c.storage = storage;
1034 175 c.store = store;
1035 {
1036 /* A node clusters into at most fan_out groups (the flat root's k =
1037 * nlist <= fan_out). The assignment scratch covers the root's full
1038 * sample; deeper slices are strictly smaller. */
1039 175 c.replay_cents = vs_alloc((size_t)c.fan_out * dim * sizeof(float));
1040 175 c.assign_scratch = vs_alloc((size_t)nvecs * sizeof(uint16_t));
1041 175 c.leaf_blk_scratch = vs_alloc((size_t)c.fan_out * sizeof(BlockNumber));
1042 }
1043 175 c.rq_params = rq_params;
1044 175 c.global_mean = global_mean;
1045 175 c.first_posting = first_posting;
1046 175 c.next_blk = first_centroid;
1047 175 c.on_leaf = on_leaf;
1048 175 c.on_leaf_arg = on_leaf_arg;
1049 175 c.collector = collector;
1050
1051 /* Same pass-scoped scratch as the PLAN pass -- the leak-containment
1052 * boundary (see the note there). on_leaf runs under it too; its
1053 * allocations must not outlive the call. */
1054 175 VsMemCtx scratch = vs_memctx_create(NULL, "vs stream write");
1055 175 VsMemCtx old_ctx = vs_memctx_switch(scratch);
1056 175 BlockNumber root = replay_node_recurse(&c, nvecs, 0);
1057 175 vs_memctx_switch(old_ctx);
1058 175 vs_memctx_delete(scratch);
1059
1/2
✓ Branch 0 taken 175 times.
✗ Branch 1 not taken.
 175 if (c.replay_cents != NULL)
1060 175 vs_free(c.replay_cents);
1061
1/2
✓ Branch 0 taken 175 times.
✗ Branch 1 not taken.
 175 if (c.assign_scratch != NULL)
1062 175 vs_free(c.assign_scratch);
1063
1/2
✓ Branch 0 taken 175 times.
✗ Branch 1 not taken.
 175 if (c.leaf_blk_scratch != NULL)
1064 175 vs_free(c.leaf_blk_scratch);
1065
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 151 times.
 175 return c.ok ? root : InvalidBlockNumber;
1066 }
1067
1068 /* Parallel build: write one worker-built (materialized) subtree's pages at
1069 * its reserved block range, mapping its leaves into the global leaf index
1070 * space -- see the header comment. */
1071 BlockNumber
1072 313 prism_routing_subtree_write(
1073 VsStorage *storage,
1074 const HKMeansResult *subtree,
1075 Dimension dim,
1076 DistanceMetric metric,
1077 uint32_t fan_out,
1078 uint8_t level_offset,
1079 PrismCentroidFormat format,
1080 const RaBitQParams *rq_params,
1081 const float *global_mean,
1082 BlockNumber first_posting,
1083 uint32_t leaf_offset,
1084 BlockNumber first_block,
1085 PrismStreamLeafCb on_leaf,
1086 void *on_leaf_arg,
1087 PrismExactCentroidCollector *collector)
1088 {
1089 313 uint32_t max_ent = prism_centroid_max_entries_fmt(dim, format);
1090
1091 /* Lay the subtree's nodes out at reserved blocks starting at first_block
1092 * (BFS: node 0 = subtree root at first_block). */
1093 313 BlockNumber *nfb = vs_alloc((size_t)subtree->nnodes * sizeof(BlockNumber));
1094 313 (void)prism_compute_centroid_layout(subtree, max_ent, first_block, nfb);
1095
1096 /* Leaf entries link to formula-derived posting heads (first_posting +
1097 * global leaf index); leaf_offset maps the subtree's local leaf indices to
1098 * the global index space. */
1099 /* Subtree node levels are subtree-relative; level_offset places them at
1100 * their absolute depth (the tree root above them is level 0). */
1101 /* Cosine: normalize the subtree's leaf centroids in place before the
1102 * leaf codes and head references are written from them (matches the
1103 * flat-tree path; k-means means drift below unit norm). */
1104
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 299 times.
 313 if (metric == DISTANCE_COSINE)
1105 {
1106 14 float *lv = hk_leaf_centroids(subtree);
1107
2/2
✓ Branch 0 taken 123 times.
✓ Branch 1 taken 14 times.
 137 for (uint32_t li = 0; li < subtree->nleaves; li++)
1108 123 vs_l2_normalize(lv + (size_t)li * dim, dim);
1109 }
1110
1111 313 prism_write_centroid_tree(
1112 storage,
1113 subtree,
1114 dim,
1115 fan_out,
1116 level_offset,
1117 format,
1118 rq_params,
1119 global_mean,
1120 first_posting + leaf_offset,
1121 nfb,
1122 NULL,
1123 collector);
1124 313 vs_free(nfb);
1125
1126 /* Head pages carry pt_centroid from the resident float leaf centroids. */
1127
1/2
✓ Branch 0 taken 313 times.
✗ Branch 1 not taken.
 313 if (on_leaf != NULL)
1128 {
1129 313 const float *leaves = hk_leaf_centroids(subtree);
1130
2/2
✓ Branch 0 taken 2284 times.
✓ Branch 1 taken 313 times.
 2597 for (uint32_t i = 0; i < subtree->nleaves; i++)
1131 2284 on_leaf(on_leaf_arg, leaf_offset + i, leaves + (size_t)i * dim);
1132 }
1133
1134 313 return first_block;
1135 }
1136
1137 /* Resolve the tree fan-out when the user left it at the default: sqrt of
1138 * the partition count, falling to cbrt when that exceeds a page's worth of
1139 * children. */
1140 uint32_t
1141 281 prism_auto_fan_out(uint32_t fan_out, uint32_t nlist, uint32_t default_fan_out)
1142 {
1143
3/4
✓ Branch 0 taken 256 times.
✓ Branch 1 taken 25 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 86 times.
 281 if (fan_out != default_fan_out || nlist <= fan_out)
1144 170 return (nlist <= fan_out) ? nlist : fan_out;
1145
1146 111 uint32_t f = (uint32_t)ceil(sqrt((double)nlist));
1147
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 111 times.
 111 if (f > 256)
1148 ✗ f = (uint32_t)ceil(cbrt((double)nlist));
1149 86 return f;
1150 }
1151
1152 /* Target entries per posting list at this dimension -- see the header. */
1153 uint32_t
1154 666 prism_target_entries_per_dim(Dimension dim, uint32_t target_pages)
1155 {
1156
2/2
✓ Branch 0 taken 559 times.
✓ Branch 1 taken 107 times.
 666 if (target_pages == 0)
1157 559 target_pages = PRISM_DEFAULT_TARGET_PAGES;
1158
1159 666 uint32_t per_page = prism_posting_page_usable() /
1160 666 PRISM_POSTING_ENTRY_SIZE(dim);
1161
1162
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 666 times.
 666 if (per_page < 1)
1163 ✗ per_page = 1;
1164
1165 666 uint32_t target = per_page * target_pages;
1166
1167 128 return target > PRISM_MIN_ENTRIES_PER_LIST ? target
1168 666 : PRISM_MIN_ENTRIES_PER_LIST;
1169 }
1170
1171 /* Automatic partition count from the (estimated) row count. */
1172 uint32_t
1173 115 prism_auto_nlist(double count, Dimension dim, uint32_t target_pages)
1174 {
1175 /*
1176 * One list per prism_target_entries_per_dim(dim) vectors -- see the header
1177 * for why that is the target. The hierarchical centroid tree keeps routing
1178 * cheap even at a high list count, so nlist scales linearly with the row
1179 * count.
1180 *
1181 * Floor it at sqrt(count): for small tables count/target collapses toward
1182 * a single list, which under-partitions and can starve the k-means build.
1183 * The linear target overtakes the sqrt floor at target^2 rows (65536 at a
1184 * 256-entry target). Back-ends clamp the result to their own nlist
1185 * ceiling.
1186 */
1187
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 95 times.
 115 double c = count > 1.0 ? count : 1.0;
1188 115 uint32_t target = prism_target_entries_per_dim(dim, target_pages);
1189 115 uint32_t linear = (uint32_t)(c / (double)target + 0.5);
1190 115 uint32_t min_lists = (uint32_t)sqrt(c);
1191 115 uint32_t nlist = linear > min_lists ? linear : min_lists;
1192
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 115 return nlist < 1 ? 1 : nlist;
1193 }
1194
1195 /*
1196 * Vectors per posting list at a given row count -- the resting size
1197 * maintenance aims each list at. Inverts prism_auto_nlist so the two cannot
1198 * drift apart; see the header for the sqrt-floor regime and the nlist == 0
1199 * convention.
1200 */
1201 uint32_t
1202 20 prism_target_entries_per_list(
1203 double count, uint32_t nlist, Dimension dim, uint32_t target_pages)
1204 {
1205
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 8 times.
 20 double c = count > 1.0 ? count : 1.0;
1206
1207
2/2
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 2 times.
 20 if (nlist == 0)
1208 18 nlist = prism_auto_nlist(c, dim, target_pages);
1209
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
 20 if (nlist < 1)
1210 ✗ nlist = 1;
1211
1212 20 uint32_t per_list = (uint32_t)(c / (double)nlist + 0.5);
1213
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 20 return per_list < 1 ? 1 : per_list;
1214 }
1215
1216 /* 2nd-nearest cluster from a distance-sorted candidate list, or
1217 * primary_cluster when the runner-up is not within epsilon (no replication
1218 * warranted) -- see the header comment. */
1219 uint32_t
1220 ✗ prism_find_secondary_cluster(
1221 const uint32_t *cand_leaves,
1222 const Distance *cand_dists,
1223 uint32_t ncand,
1224 uint32_t primary_cluster,
1225 Distance primary_dist,
1226 double epsilon)
1227 {
1228 /* Candidates are sorted by ascending distance, so the first one
1229 * that is not the primary is the 2nd-nearest centroid. */
1230 ✗ for (uint32_t i = 0; i < ncand; i++)
1231 {
1232 ✗ if (cand_leaves[i] == primary_cluster)
1233 ✗ continue;
1234
1235 ✗ Distance best2 = cand_dists[i];
1236 ✗ double gap = (double)best2 - (double)primary_dist;
1237 ✗ double gap_ratio = (primary_dist != 0.0)
1238 ✗ ? gap / fabs((double)primary_dist)
1239 ✗ : INFINITY;
1240
1241 ✗ if (gap_ratio <= epsilon)
1242 ✗ return cand_leaves[i];
1243 ✗ break;
1244 }
1245
1246 ✗ return primary_cluster;
1247 }
1248
1249 /*
1250 * SOAR secondary cluster: the leaf (other than the primary) minimizing the
1251 * orthogonality-amplified distance, decomposed so both terms use the SIMD
1252 * vecops kernels:
1253 *
1254 * oa(c) = ||v - c||^2 + lambda * (r_hat . (v - c))^2
1255 * = ||v - c||^2 + lambda * (r_hat.v - r_hat.c)^2
1256 *
1257 * r_hat.v is constant across centroids, so only ||v - c||^2 and r_hat.c are
1258 * per-centroid. Since lambda * (...)^2 >= 0, ||v - c||^2 is a lower bound on
1259 * oa: when it already exceeds the running best we skip the dot product (exact
1260 * pruning — no recall impact).
1261 *
1262 * The search set is `count` leaves: the ids cand_leaves[0..count) when
1263 * cand_leaves is non-NULL, otherwise leaves 0..count (a full scan, called with
1264 * count == nleaves). The candidate form is O(count) rather than O(nleaves);
1265 * because the minimizer always has a small ||v - c||^2 (a far leaf cannot win)
1266 * it lies among the nearest leaves the beam already found, so the result is
1267 * unchanged while scaling to large nlist — where a full scan would read the
1268 * entire (multi-hundred-MB) leaf-centroid array per replicated vector.
1269 */
1270 uint32_t
1271 35492 prism_find_soar_secondary(
1272 const float *vec,
1273 const float *leaf_centroids,
1274 const uint32_t *cand_leaves,
1275 uint32_t count,
1276 Dimension dim,
1277 uint32_t primary_cluster,
1278 const float *normalized_residual,
1279 double lambda)
1280 {
1281 35492 float qrv = vs_dot_product(normalized_residual, vec, dim);
1282 35492 float lam = (float)lambda;
1283 35492 float best_oa = INFINITY;
1284 35492 uint32_t best_c = primary_cluster;
1285
1286
2/2
✓ Branch 0 taken 1287397 times.
✓ Branch 1 taken 35492 times.
 1322889 for (uint32_t k = 0; k < count; k++)
1287 {
1288
2/2
✓ Branch 0 taken 124806 times.
✓ Branch 1 taken 1162591 times.
 1287397 uint32_t i = cand_leaves ? cand_leaves[k] : k;
1289
2/2
✓ Branch 0 taken 34944 times.
✓ Branch 1 taken 1252453 times.
 1287397 if (i == primary_cluster)
1290 34944 continue;
1291
1292 1252453 const float *cent = leaf_centroids + (size_t)i * dim;
1293
1294 1252453 float l2 = vs_l2_distance_squared(vec, cent, dim);
1295
2/2
✓ Branch 0 taken 1168150 times.
✓ Branch 1 taken 84303 times.
 1252453 if (l2 >= best_oa)
1296 1168150 continue; /* oa >= l2 >= best_oa: cannot improve */
1297
1298 84303 float rc = vs_dot_product(normalized_residual, cent, dim);
1299 84303 float gap = qrv - rc;
1300 84303 float oa = l2 + lam * gap * gap;
1301
2/2
✓ Branch 0 taken 39392 times.
✓ Branch 1 taken 44911 times.
 84303 if (oa < best_oa)
1302 {
1303 1051353 best_oa = oa;
1304 1051353 best_c = i;
1305 }
1306 }
1307
1308 35492 return best_c;
1309 }
1310
1311 /* Log the per-phase build time summary. */
1312 void
1313 194 prism_build_stats_print(const PrismBuildStats *s)
1314 {
1315 /* Per-phase breakdown — the single shared build summary, filled by both
1316 * the PostgreSQL build (via the build-progress seam) and the standalone
1317 * build. */
1318
1/2
✓ Branch 1 taken 194 times.
✗ Branch 2 not taken.
 194 vs_log("build: sample %.1fms, kmeans %.1fms, refine %.1fms, setup "
1319 "%.1fms, "
1320 "posting %.1fms, centroid %.1fms, total %.1fms\n",
1321 s->ms_sample,
1322 s->ms_kmeans,
1323 s->ms_refine,
1324 s->ms_setup,
1325 s->ms_posting,
1326 s->ms_centroid,
1327 s->ms_total);
1328
1329 /* Posting-merge sub-detail: only the standalone parallel path populates
1330 * these. Skip the line (and its zeros) when unset — e.g. the PostgreSQL
1331 * build, which tracks posting as one phase. */
1332
3/6
✓ Branch 0 taken 194 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 194 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 194 times.
 194 if (s->ms_parallel > 0.0 || s->ms_merge > 0.0 || s->total_pages > 0)
1333 ✗ vs_log("build: posting detail — parallel %.1fms + merge %.1fms, "
1334 "%u workers + leader, %u pages, "
1335 "%u partial pages merged into %u\n",
1336 s->ms_parallel,
1337 s->ms_merge,
1338 s->nworkers,
1339 s->total_pages,
1340 s->merge_input,
1341 s->merge_output);
1342 194 }
1343