| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | /* | ||
| 2 | * Copyright (c) 2026 Tiger Data, Inc. | ||
| 3 | * Licensed under the PostgreSQL License. See LICENSE for details. | ||
| 4 | * | ||
| 5 | * posting_build.h - Posting list construction | ||
| 6 | * | ||
| 7 | * Vector-to-cluster assignment (tree descent, SOAR, boundary | ||
| 8 | * replication) and streaming page builder. Used by both standalone | ||
| 9 | * and PostgreSQL builds. No PostgreSQL dependencies. | ||
| 10 | */ | ||
| 11 | |||
| 12 | #ifndef PRISM_POSTING_BUILD_H | ||
| 13 | #define PRISM_POSTING_BUILD_H | ||
| 14 | |||
| 15 | #include <stdbool.h> | ||
| 16 | #include <stdint.h> | ||
| 17 | |||
| 18 | #include "algo/hkmeans.h" | ||
| 19 | #include "core/atomics.h" | ||
| 20 | #include "core/memory.h" | ||
| 21 | #include "core/types.h" | ||
| 22 | #include "index/posting_page.h" | ||
| 23 | #include "index/query_scan.h" | ||
| 24 | #include "index/storage.h" | ||
| 25 | #include "quant/rabitq.h" | ||
| 26 | |||
| 27 | /* Opaque: the cluster-keyed sorter (defined in parallel_build.h). */ | ||
| 28 | typedef struct PrismSorter PrismSorter; | ||
| 29 | |||
| 30 | /* ---------------------------------------------------------------- | ||
| 31 | * Vector-to-cluster assignment | ||
| 32 | * ---------------------------------------------------------------- */ | ||
| 33 | |||
| 34 | #define PRISM_INVALID_CLUSTER UINT32_MAX | ||
| 35 | |||
| 36 | typedef struct PrismAssignParams | ||
| 37 | { | ||
| 38 | Dimension dim; | ||
| 39 | DistanceMetric metric; | ||
| 40 | double soar_lambda; | ||
| 41 | double boundary_epsilon; | ||
| 42 | } PrismAssignParams; | ||
| 43 | |||
| 44 | typedef struct PrismBuildAssignment | ||
| 45 | { | ||
| 46 | uint32_t primary; | ||
| 47 | uint32_t secondary; | ||
| 48 | const float *enc_vector; | ||
| 49 | } PrismBuildAssignment; | ||
| 50 | |||
| 51 | /* | ||
| 52 | * Beam parameters of the build-time leaf-candidate pool. The page-backed | ||
| 53 | * assign path routes each row with nprobe = PRISM_SECONDARY_TOPK, then | ||
| 54 | * exact-re-ranks those leaf finalists against their head pages' | ||
| 55 | * full-precision pt_centroids: the primary is the exact-nearest finalist, | ||
| 56 | * the boundary cluster the exact 2nd-nearest, and SOAR's | ||
| 57 | * orthogonality-amplified search scans all of them exactly (its optimum | ||
| 58 | * need not be the nearest-by-distance leaf, but it is always among the | ||
| 59 | * near ones). The in-RAM assign path (prism_build_assign_vector) descends | ||
| 60 | * the exact float tree with the same widths. | ||
| 61 | * | ||
| 62 | * TOPK 32 / BEAM_WIDTH 64: with the internal levels scored exactly (see | ||
| 63 | * PRISM_BUILD_CENTROID_BEAM_SCALE below), leaf assignment quality is bounded | ||
| 64 | * by how often the true nearest leaf sits inside the estimated top-TOPK of | ||
| 65 | * the exactly-chosen parents' children; widening 8/16 -> 32/64 measurably | ||
| 66 | * recovers assignment recall at acceptable build cost. | ||
| 67 | */ | ||
| 68 | #define PRISM_SECONDARY_TOPK 32 | ||
| 69 | #define PRISM_SECONDARY_BEAM_WIDTH 64 | ||
| 70 | |||
| 71 | /* | ||
| 72 | * Build-time centroid routing accuracy. | ||
| 73 | * | ||
| 74 | * The build assigns every vector exactly once, so it routes the centroid | ||
| 75 | * descent for accuracy, not query speed. It must NOT inherit the | ||
| 76 | * query-tuned prism.centroid_beam_scale, which trades recall for QPS: a | ||
| 77 | * small beam_scale makes the beam narrower than the PRISM_SECONDARY_TOPK | ||
| 78 | * candidates the secondary (boundary/SOAR) search requests, loosening | ||
| 79 | * clustering and permanently lowering query recall. | ||
| 80 | * | ||
| 81 | * BEAM_SCALE = 1.0 makes beam_w = nprobe = PRISM_SECONDARY_TOPK, i.e. the | ||
| 82 | * descent keeps as many candidates as the secondary search returns. This | ||
| 83 | * matches the build query-state buffers (beam_results / centroid_scratch | ||
| 84 | * are sized for max_nprobe = PRISM_SECONDARY_TOPK). The query-tuned default | ||
| 85 | * (0.25) instead yields beam_w = 2, far narrower than the k candidates the | ||
| 86 | * descent must produce. | ||
| 87 | * | ||
| 88 | * 1.0 needs no slack above that: the build descent scores the INTERNAL | ||
| 89 | * tree levels against the exact float centroids collected while the tree | ||
| 90 | * streamed to pages (PrismExactInternalCentroids), so the kept parents are the | ||
| 91 | * true top-beam_w — there is no estimate noise at those levels for a wider | ||
| 92 | * beam to paper over. (Widening was only ever a half-measure for that | ||
| 93 | * noise: beam scale 4 recovered about half the misassignment loss at 4x | ||
| 94 | * the descent cost; exact internal scoring removes the cause.) The leaf | ||
| 95 | * level keeps the 1-bit estimates — it is far too large to collect — and | ||
| 96 | * its noise is absorbed by the exact re-rank of the TOPK finalists above. | ||
| 97 | * | ||
| 98 | * ERROR_SCALE = 0 matches the query default; with exact internal scoring | ||
| 99 | * the bounds are exact (error = 0) and pruning slack is meaningless. It | ||
| 100 | * must stay 0 here — a positive error_scale in the single-candidate refine | ||
| 101 | * route (nprobe = 1) overflows the beam-search candidate buffer and | ||
| 102 | * corrupts the build. | ||
| 103 | */ | ||
| 104 | #define PRISM_BUILD_CENTROID_BEAM_SCALE 1.0f | ||
| 105 | #define PRISM_BUILD_CENTROID_ERROR_SCALE 0.0f | ||
| 106 | |||
| 107 | typedef struct PrismBuildWorkerBufs | ||
| 108 | { | ||
| 109 | float *norm_buf; /* [dim] for cosine normalization */ | ||
| 110 | float *residual_buf; /* [dim] for SOAR residual computation */ | ||
| 111 | uint32_t *cand_leaves; /* [PRISM_SECONDARY_TOPK] beam candidates */ | ||
| 112 | Distance *cand_dists; /* [PRISM_SECONDARY_TOPK] candidate distances */ | ||
| 113 | } PrismBuildWorkerBufs; | ||
| 114 | |||
| 115 | PrismBuildWorkerBufs prism_build_worker_bufs_create(Dimension dim); | ||
| 116 | void prism_build_worker_bufs_free(PrismBuildWorkerBufs *bufs); | ||
| 117 | |||
| 118 | PrismBuildAssignment prism_build_assign_vector( | ||
| 119 | const HKMeansResult *tree, | ||
| 120 | const float *vec, | ||
| 121 | const PrismAssignParams *params, | ||
| 122 | PrismBuildWorkerBufs *bufs); | ||
| 123 | |||
| 124 | /* ---------------------------------------------------------------- | ||
| 125 | * Page-backed build routing (unified with the query/insert path) | ||
| 126 | * | ||
| 127 | * Routes each vector to its posting list exactly as a query does -- | ||
| 128 | * prism_query_route over the centroid pages -- then exact-re-ranks the TOPK | ||
| 129 | * leaf finalists against their head pages' full-precision pt_centroids | ||
| 130 | * (the primary is the exact-nearest finalist), encodes the RaBitQ | ||
| 131 | * residual against the chosen list's pt_centroid and streams it to the | ||
| 132 | * cluster-keyed sorter (primary + optional SOAR / boundary secondary). | ||
| 133 | * Shared by the serial build and the parallel posting workers. The | ||
| 134 | * descent is the query's; on top of it the build exact-re-ranks the | ||
| 135 | * finalists (queries re-rank their probe set instead, inserts keep the | ||
| 136 | * plain estimate route). The centroid and head pages must already be | ||
| 137 | * written when this runs. | ||
| 138 | * ---------------------------------------------------------------- */ | ||
| 139 | typedef struct PrismBuildRouteCtx | ||
| 140 | { | ||
| 141 | PrismQueryState *qs; /* routing state (not owned) */ | ||
| 142 | PrismSorter *sorter; /* cluster-keyed output (not owned) */ | ||
| 143 | const RaBitQParams *rq_params; /* not owned */ | ||
| 144 | VsStorage *storage; /* head-page reads (not owned) */ | ||
| 145 | BlockNumber first_posting; /* leaf c's head = first_posting + c */ | ||
| 146 | Dimension dim; | ||
| 147 | double soar_lambda; | ||
| 148 | double boundary_epsilon; | ||
| 149 | |||
| 150 | /* Owned scratch (allocated in init, freed in cleanup). */ | ||
| 151 | float *cand_pt; /* [PRISM_SECONDARY_TOPK * dim] gathered pt_centroids */ | ||
| 152 | uint32_t *cand_leaf; /* [PRISM_SECONDARY_TOPK] leaf per candidate */ | ||
| 153 | Distance *cand_dist; /* [PRISM_SECONDARY_TOPK] candidate distances */ | ||
| 154 | float *pt_r; /* [dim] rotated residual scratch */ | ||
| 155 | RaBitQData *enc_buf; /* RaBitQ encode output */ | ||
| 156 | RaBitQScratch enc_scratch; | ||
| 157 | char *entry; /* [prism_posting_entry_size(dim)] */ | ||
| 158 | |||
| 159 | /* Counters. */ | ||
| 160 | double indtuples; | ||
| 161 | double soar_dupes; | ||
| 162 | } PrismBuildRouteCtx; | ||
| 163 | |||
| 164 | void prism_build_route_ctx_init( | ||
| 165 | PrismBuildRouteCtx *ctx, | ||
| 166 | PrismQueryState *qs, | ||
| 167 | PrismSorter *sorter, | ||
| 168 | const RaBitQParams *rq_params, | ||
| 169 | VsStorage *storage, | ||
| 170 | BlockNumber first_posting, | ||
| 171 | Dimension dim, | ||
| 172 | double soar_lambda, | ||
| 173 | double boundary_epsilon); | ||
| 174 | |||
| 175 | void prism_build_route_ctx_cleanup(PrismBuildRouteCtx *ctx); | ||
| 176 | |||
| 177 | /* Route one vector, encode, and stream its entries to the sorter. Returns | ||
| 178 | * true when a secondary (SOAR / boundary) replica was also emitted. */ | ||
| 179 | bool prism_build_route_emit( | ||
| 180 | PrismBuildRouteCtx *ctx, const float *vec, ItemPointerData tid); | ||
| 181 | |||
| 182 | /* Map a routed posting-head block back to its leaf index. Head blocks are the | ||
| 183 | * formula first_posting + leaf, so this is a subtraction. Shared by the | ||
| 184 | * route/encode path and the page-backed refine pass. */ | ||
| 185 | static inline uint32_t | ||
| 186 | 13367066 | prism_route_head_to_leaf(BlockNumber first_posting, BlockNumber head) | |
| 187 | { | ||
| 188 |
2/2✓ Branch 0 taken 48600 times.
✓ Branch 1 taken 24000 times.
|
13367066 | return (uint32_t)(head - first_posting); |
| 189 | } | ||
| 190 | |||
| 191 | /* | ||
| 192 | * Shared refine kernels: the serial and parallel refine passes route, | ||
| 193 | * filter, and average identically -- only the accumulator ownership | ||
| 194 | * (private vs striped-locked DSM) differs, so that add stays with the | ||
| 195 | * caller. | ||
| 196 | * | ||
| 197 | * prism_refine_route_row routes one row page-backed (k=1, exactly as the | ||
| 198 | * query and insert paths do), maps the head back to its leaf, and returns | ||
| 199 | * the vector to accumulate (the normalized copy in scratch for cosine) -- | ||
| 200 | * or NULL when the row routed nowhere or outside [tile_lo, tile_hi). | ||
| 201 | * *out_idx is the tile-relative leaf index. | ||
| 202 | */ | ||
| 203 | const float *prism_refine_route_row( | ||
| 204 | struct PrismQueryState *qs, | ||
| 205 | BlockNumber first_posting, | ||
| 206 | const float *vec, | ||
| 207 | Dimension dim, | ||
| 208 | bool cosine, | ||
| 209 | float *scratch, | ||
| 210 | uint32_t tile_lo, | ||
| 211 | uint32_t tile_hi, | ||
| 212 | uint32_t *out_idx); | ||
| 213 | |||
| 214 | /* | ||
| 215 | * Divide one tile's sums by their counts and hand each refined leaf mean to | ||
| 216 | * write_head (leaves with no routed rows keep their sample-trained head). | ||
| 217 | * scratch is a caller-owned [dim] float buffer. | ||
| 218 | */ | ||
| 219 | typedef void (*PrismLeafWriteFn)(void *ctx, uint32_t leaf, const float *vec); | ||
| 220 | |||
| 221 | /* | ||
| 222 | * Shared leaf-head writer: writes leaf's posting-list head page (at the | ||
| 223 | * formula-derived block first_posting + leaf) carrying pt_centroid = | ||
| 224 | * P^T * centroid -- the encode reference the scan reads. Used as the | ||
| 225 | * PrismLeafWriteFn of both the streaming tree write and the refine pass, by | ||
| 226 | * the serial build and the parallel leader alike. pt is caller-owned [dim] | ||
| 227 | * scratch. | ||
| 228 | */ | ||
| 229 | typedef struct PrismHeadWriteCtx | ||
| 230 | { | ||
| 231 | VsStorage *storage; | ||
| 232 | const RaBitQParams *rq_params; | ||
| 233 | Dimension dim; | ||
| 234 | bool fastscan; | ||
| 235 | BlockNumber first_posting; /* leaf c's head = first_posting + c */ | ||
| 236 | float *pt; /* [dim] scratch */ | ||
| 237 | } PrismHeadWriteCtx; | ||
| 238 | |||
| 239 | /* Stack-init/cleanup pair: allocates/frees the [dim] pt scratch. */ | ||
| 240 | static inline void | ||
| 241 | 281 | prism_head_write_ctx_init( | |
| 242 | PrismHeadWriteCtx *h, | ||
| 243 | VsStorage *storage, | ||
| 244 | const RaBitQParams *rq_params, | ||
| 245 | Dimension dim, | ||
| 246 | bool fastscan, | ||
| 247 | BlockNumber first_posting) | ||
| 248 | { | ||
| 249 | 281 | h->storage = storage; | |
| 250 | 281 | h->rq_params = rq_params; | |
| 251 | 281 | h->dim = dim; | |
| 252 | 281 | h->fastscan = fastscan; | |
| 253 | 281 | h->first_posting = first_posting; | |
| 254 | 281 | h->pt = vs_alloc((size_t)dim * sizeof(float)); | |
| 255 | 82 | } | |
| 256 | |||
| 257 | static inline void | ||
| 258 | 281 | prism_head_write_ctx_cleanup(PrismHeadWriteCtx *h) | |
| 259 | { | ||
| 260 | 281 | vs_free(h->pt); | |
| 261 | 82 | } | |
| 262 | |||
| 263 | void prism_write_leaf_head(void *arg, uint32_t leaf, const float *centroid); | ||
| 264 | |||
| 265 | /* | ||
| 266 | * Build-time page-backed router base: the same PrismIndexBase the query and | ||
| 267 | * insert paths build, so the build scan routes each row identically. Both | ||
| 268 | * back-ends must construct it the same way -- this is the one place. The | ||
| 269 | * caller owns pt_global_mean ([dim], filled with the rotated global mean) | ||
| 270 | * and frees it after prism_query_state_cleanup. | ||
| 271 | */ | ||
| 272 | static inline void | ||
| 273 | 341 | prism_build_router_base_init( | |
| 274 | PrismIndexBase *base, | ||
| 275 | RaBitQParams *params, | ||
| 276 | VsStorage *storage, | ||
| 277 | Dimension dim, | ||
| 278 | uint8_t nlevels, | ||
| 279 | BlockNumber first_centroid, | ||
| 280 | DistanceMetric metric, | ||
| 281 | PrismCentroidFormat centroid_format, | ||
| 282 | int fastscan_bits, | ||
| 283 | double error_scale, | ||
| 284 | double beam_scale, | ||
| 285 | uint32_t fan_out, | ||
| 286 | uint32_t nlist, | ||
| 287 | uint64_t rabitq_seed, | ||
| 288 | const float *global_mean, | ||
| 289 | float *pt_global_mean) | ||
| 290 | { | ||
| 291 | 341 | memset(base, 0, sizeof(*base)); | |
| 292 | 341 | base->params = params; | |
| 293 | 341 | base->pt_global_mean = pt_global_mean; | |
| 294 | 341 | vs_rabitq_rotate(params, global_mean, base->pt_global_mean); | |
| 295 | 341 | base->rabitq_seed = rabitq_seed; | |
| 296 | 341 | base->centroid_storage = storage; | |
| 297 | 341 | base->posting_storage = storage; | |
| 298 | 341 | base->page_base = NULL; | |
| 299 | 341 | base->dim = dim; | |
| 300 | 341 | base->nlevels = nlevels; | |
| 301 | 341 | base->first_centroid = first_centroid; | |
| 302 | 341 | base->metric = metric; | |
| 303 | 341 | base->centroid_format = centroid_format; | |
| 304 | 576 | base->fastscan = (centroid_format == PRISM_CENTROID_FMT_FASTSCAN) | |
| 305 | ? fastscan_bits | ||
| 306 |
2/2✓ Branch 0 taken 57 times.
✓ Branch 1 taken 284 times.
|
341 | : 0; |
| 307 | 341 | base->centroid_error_scale = error_scale; | |
| 308 | 341 | base->centroid_beam_scale = beam_scale; | |
| 309 |
1/2✓ Branch 0 taken 106 times.
✗ Branch 1 not taken.
|
341 | base->fan_out = (uint8_t)(fan_out <= UINT8_MAX ? fan_out : UINT8_MAX); |
| 310 | 341 | base->nlist = nlist; | |
| 311 | 341 | } | |
| 312 | |||
| 313 | void prism_refine_write_means( | ||
| 314 | const double *sums, | ||
| 315 | const uint64_t *counts, | ||
| 316 | uint32_t lo, | ||
| 317 | uint32_t hi, | ||
| 318 | Dimension dim, | ||
| 319 | float *scratch, | ||
| 320 | PrismLeafWriteFn write_head, | ||
| 321 | void *write_head_ctx); | ||
| 322 | |||
| 323 | /* ---------------------------------------------------------------- | ||
| 324 | * Page format ops — the only part that differs between formats | ||
| 325 | * | ||
| 326 | * write_entry: add one encoded entry to the in-memory page. | ||
| 327 | * Returns false if the page is full and needs flushing first. | ||
| 328 | * reinit_page: re-initialize the page buffer after a flush. | ||
| 329 | * finalize: called before the last flush (e.g., write partial | ||
| 330 | * fastscan group). | ||
| 331 | * cleanup: free format-specific scratch buffers. | ||
| 332 | * ---------------------------------------------------------------- */ | ||
| 333 | |||
| 334 | struct PrismPostingBuilder; | ||
| 335 | |||
| 336 | typedef struct PrismPostingPageOps | ||
| 337 | { | ||
| 338 | bool (*write_entry)( | ||
| 339 | struct PrismPostingBuilder *b, | ||
| 340 | ItemPointerData tid, | ||
| 341 | float f_add, | ||
| 342 | float f_rescale, | ||
| 343 | float f_error, | ||
| 344 | const uint8_t *bits); | ||
| 345 | void (*reinit_page)(struct PrismPostingBuilder *b); | ||
| 346 | void (*finalize)(struct PrismPostingBuilder *b); | ||
| 347 | void (*cleanup)(struct PrismPostingBuilder *b); | ||
| 348 | } PrismPostingPageOps; | ||
| 349 | |||
| 350 | /* Fastscan group staging buffer */ | ||
| 351 | typedef struct FsGroupStage | ||
| 352 | { | ||
| 353 | ItemPointerData tids[VS_FASTSCAN_GROUP]; | ||
| 354 | float f_add[VS_FASTSCAN_GROUP]; | ||
| 355 | float f_rescale[VS_FASTSCAN_GROUP]; | ||
| 356 | float f_error[VS_FASTSCAN_GROUP]; | ||
| 357 | uint32_t count; | ||
| 358 | } FsGroupStage; | ||
| 359 | |||
| 360 | /* ---------------------------------------------------------------- | ||
| 361 | * Unified builder state | ||
| 362 | * ---------------------------------------------------------------- */ | ||
| 363 | |||
| 364 | typedef struct PrismPostingBuilder | ||
| 365 | { | ||
| 366 | /* Common fields */ | ||
| 367 | VsStorage *storage; | ||
| 368 | const RaBitQParams *params; | ||
| 369 | Dimension dim; | ||
| 370 | uint32_t cluster_id; | ||
| 371 | const float *centroid; | ||
| 372 | |||
| 373 | BlockNumber head_blkno; | ||
| 374 | BlockNumber prev_blkno; | ||
| 375 | bool is_first; | ||
| 376 | |||
| 377 | /* | ||
| 378 | * Head-metadata bookkeeping. owns_head is set for head builders (those | ||
| 379 | * that produce the FIRST page) and drives whether finish() stamps the | ||
| 380 | * head's live_count / tail_blkno. n_entries counts every entry added | ||
| 381 | * through this builder, so a head builder that holds the whole chain | ||
| 382 | * (serial build, and the parallel leader's head builder when no | ||
| 383 | * continuations were streamed) stamps an exact live_count with no walk. | ||
| 384 | */ | ||
| 385 | bool owns_head; | ||
| 386 | uint32_t n_entries; | ||
| 387 | |||
| 388 | char mem_page[BLCKSZ] __attribute__((aligned(8))); | ||
| 389 | bool page_dirty; | ||
| 390 | /* Head page was adopted from storage (see | ||
| 391 | * prism_posting_builder_adopt_head): when nothing is appended, the | ||
| 392 | * on-disk head is already final and is neither rewritten nor | ||
| 393 | * restamped. */ | ||
| 394 | bool adopted_head; | ||
| 395 | |||
| 396 | BlockNumber fixed_first_blkno; | ||
| 397 | |||
| 398 | RaBitQData *enc_buf; | ||
| 399 | RaBitQScratch enc_scratch; | ||
| 400 | |||
| 401 | /* Page format dispatch */ | ||
| 402 | const PrismPostingPageOps *page_ops; | ||
| 403 | |||
| 404 | /* Format-specific state (only fastscan uses this) */ | ||
| 405 | struct | ||
| 406 | { | ||
| 407 | FsGroupStage grp; | ||
| 408 | uint32_t groups_on_page; | ||
| 409 | uint32_t max_groups; | ||
| 410 | uint8_t *bits_buf; | ||
| 411 | uint8_t *codes_buf; | ||
| 412 | } fs; | ||
| 413 | } PrismPostingBuilder; | ||
| 414 | |||
| 415 | /* | ||
| 416 | * Initialize for AoS page format. Encodes vectors with RaBitQ. | ||
| 417 | */ | ||
| 418 | void prism_posting_builder_init( | ||
| 419 | PrismPostingBuilder *builder, | ||
| 420 | VsStorage *storage, | ||
| 421 | const RaBitQParams *params, | ||
| 422 | Dimension dim, | ||
| 423 | uint32_t cluster_id, | ||
| 424 | const float *centroid, | ||
| 425 | const float *pt_centroid); | ||
| 426 | |||
| 427 | /* | ||
| 428 | * Initialize for fastscan page format. When params and centroid | ||
| 429 | * are non-NULL, use _add() to encode raw vectors. When both are | ||
| 430 | * NULL, use _add_encoded() for pre-encoded data (AoS conversion). | ||
| 431 | */ | ||
| 432 | void prism_posting_builder_init_fastscan( | ||
| 433 | PrismPostingBuilder *builder, | ||
| 434 | VsStorage *storage, | ||
| 435 | const RaBitQParams *params, | ||
| 436 | Dimension dim, | ||
| 437 | uint32_t cluster_id, | ||
| 438 | const float *centroid, | ||
| 439 | const float *pt_centroid); | ||
| 440 | |||
| 441 | /* | ||
| 442 | * Format-dispatching init: fastscan page format when fastscan is true, | ||
| 443 | * AoS otherwise. Same arguments as the two inits above. | ||
| 444 | */ | ||
| 445 | void prism_posting_builder_init_fmt( | ||
| 446 | PrismPostingBuilder *builder, | ||
| 447 | VsStorage *storage, | ||
| 448 | const RaBitQParams *params, | ||
| 449 | Dimension dim, | ||
| 450 | uint32_t cluster_id, | ||
| 451 | const float *centroid, | ||
| 452 | const float *pt_centroid, | ||
| 453 | bool fastscan); | ||
| 454 | |||
| 455 | /* | ||
| 456 | * Adopt the pre-written (empty) posting-list head at head_blk and append | ||
| 457 | * into it. The page-backed build writes each head before the posting scan | ||
| 458 | * (the scan routes and encodes against its pt_centroid); adopting appends | ||
| 459 | * to that same page instead of constructing a second head that must agree | ||
| 460 | * with it, and a cluster that receives no entries keeps its on-disk head | ||
| 461 | * untouched. Entries must arrive pre-encoded (prism_posting_entry_add). | ||
| 462 | */ | ||
| 463 | void prism_posting_builder_adopt_head( | ||
| 464 | PrismPostingBuilder *builder, | ||
| 465 | VsStorage *storage, | ||
| 466 | const RaBitQParams *params, | ||
| 467 | Dimension dim, | ||
| 468 | uint32_t cluster_id, | ||
| 469 | BlockNumber head_blk, | ||
| 470 | bool fastscan); | ||
| 471 | |||
| 472 | /* | ||
| 473 | * Pin the first page to a specific block number. The first flush | ||
| 474 | * writes to this block; later pages are appended via new_page. | ||
| 475 | */ | ||
| 476 | void prism_posting_builder_set_first_blkno( | ||
| 477 | PrismPostingBuilder *builder, BlockNumber blkno); | ||
| 478 | |||
| 479 | /* | ||
| 480 | * Derive RaBitQ error bound from encoding factors. | ||
| 481 | */ | ||
| 482 | |||
| 483 | /* | ||
| 484 | * Add a raw vector. Encodes with RaBitQ relative to centroid. | ||
| 485 | */ | ||
| 486 | void prism_posting_builder_add( | ||
| 487 | PrismPostingBuilder *builder, | ||
| 488 | ItemPointerData tid, | ||
| 489 | const float *vector); | ||
| 490 | |||
| 491 | /* | ||
| 492 | * Add a raw vector, optionally stamping it unreachable: estimated distance | ||
| 493 | * +inf with zero error, so every scan prunes it before it can enter the top-k | ||
| 494 | * threshold heap. For vectors with no defined distance under the index metric | ||
| 495 | * (a zero-norm vector under cosine) -- see mark_entry_unreachable in | ||
| 496 | * posting_build.c for why a naive encoding is actively harmful. The row stays | ||
| 497 | * indexed; it just can never be a result. | ||
| 498 | */ | ||
| 499 | void prism_posting_builder_add_ex( | ||
| 500 | PrismPostingBuilder *builder, | ||
| 501 | ItemPointerData tid, | ||
| 502 | const float *vector, | ||
| 503 | bool unreachable); | ||
| 504 | |||
| 505 | /* | ||
| 506 | * Add a pre-encoded entry. No RaBitQ encoding — data is already | ||
| 507 | * quantized. Works with both AoS and fastscan formats. | ||
| 508 | */ | ||
| 509 | void prism_posting_builder_add_encoded( | ||
| 510 | PrismPostingBuilder *builder, | ||
| 511 | ItemPointerData tid, | ||
| 512 | float f_add, | ||
| 513 | float f_rescale, | ||
| 514 | float f_error, | ||
| 515 | const uint8_t *bits); | ||
| 516 | |||
| 517 | BlockNumber prism_posting_builder_finish(PrismPostingBuilder *builder); | ||
| 518 | |||
| 519 | void prism_posting_builder_cleanup(PrismPostingBuilder *builder); | ||
| 520 | |||
| 521 | /* ---------------------------------------------------------------- | ||
| 522 | * Compact posting entry for the cluster-sorted build path. | ||
| 523 | * | ||
| 524 | * A fixed-size blob keyed (externally) by cluster id, carrying the RaBitQ code | ||
| 525 | * so the sort moves ~108 bytes/vector instead of dim*4. Layout: | ||
| 526 | * [ItemPointerData tid][float f_add][float f_rescale][float f_error] | ||
| 527 | * [uint8 sign_bits[(dim+7)/8]] | ||
| 528 | * prism_posting_entry_encode() fills it (relative to the assigned cluster | ||
| 529 | * centroid); prism_posting_entry_add() replays it into a builder via | ||
| 530 | * add_encoded. Used by both the serial (build.c) and parallel | ||
| 531 | * (sort-seam) builds, so the format has a single definition. | ||
| 532 | * ---------------------------------------------------------------- */ | ||
| 533 | Size prism_posting_entry_size(Dimension dim); | ||
| 534 | |||
| 535 | void prism_posting_entry_encode( | ||
| 536 | const RaBitQParams *params, | ||
| 537 | const float *vec, | ||
| 538 | const float *centroid, | ||
| 539 | Dimension dim, | ||
| 540 | RaBitQData *enc_buf, /* scratch, VS_RABITQ_DATA_SIZE(dim) */ | ||
| 541 | RaBitQScratch *scratch, /* scratch, scratch_init(dim) */ | ||
| 542 | ItemPointerData tid, | ||
| 543 | void *out_entry); /* prism_posting_entry_size(dim) bytes */ | ||
| 544 | |||
| 545 | /* | ||
| 546 | * Page-backed twin of prism_posting_entry_encode: encodes from a pre-computed | ||
| 547 | * rotated residual (pt_query - pt_centroid) instead of a float vector + | ||
| 548 | * centroid. Used by the page-backed build, where pt_centroid is read from the | ||
| 549 | * posting head page (no in-RAM float centroids). | ||
| 550 | */ | ||
| 551 | void prism_posting_entry_encode_from_pt( | ||
| 552 | const RaBitQParams *params, | ||
| 553 | const float *pt_residual, | ||
| 554 | Dimension dim, | ||
| 555 | RaBitQData *enc_buf, | ||
| 556 | RaBitQScratch *scratch, | ||
| 557 | ItemPointerData tid, | ||
| 558 | void *out_entry); | ||
| 559 | |||
| 560 | void prism_posting_entry_add( | ||
| 561 | PrismPostingBuilder *builder, const void *entry, Dimension dim); | ||
| 562 | |||
| 563 | /* ---------------------------------------------------------------- | ||
| 564 | * Flat builder — one buffer per cluster (standalone benchmark) | ||
| 565 | * ---------------------------------------------------------------- */ | ||
| 566 | |||
| 567 | typedef struct PrismFlatPostingBuilder | ||
| 568 | { | ||
| 569 | const RaBitQParams *params; | ||
| 570 | Dimension dim; | ||
| 571 | uint32_t cluster_id; | ||
| 572 | const float *centroid; | ||
| 573 | |||
| 574 | char *buf; | ||
| 575 | uint32_t max_entries; | ||
| 576 | |||
| 577 | RaBitQData *enc_buf; | ||
| 578 | } PrismFlatPostingBuilder; | ||
| 579 | |||
| 580 | void prism_flat_posting_builder_init( | ||
| 581 | PrismFlatPostingBuilder *builder, | ||
| 582 | const RaBitQParams *params, | ||
| 583 | Dimension dim, | ||
| 584 | uint32_t cluster_id, | ||
| 585 | const float *centroid, | ||
| 586 | uint32_t count); | ||
| 587 | |||
| 588 | void prism_flat_posting_builder_add( | ||
| 589 | PrismFlatPostingBuilder *builder, | ||
| 590 | ItemPointerData tid, | ||
| 591 | const float *vector); | ||
| 592 | |||
| 593 | char *prism_flat_posting_builder_finish(PrismFlatPostingBuilder *builder); | ||
| 594 | |||
| 595 | void prism_flat_posting_builder_cleanup(PrismFlatPostingBuilder *builder); | ||
| 596 | |||
| 597 | #endif /* PRISM_POSTING_BUILD_H */ | ||
| 598 |