GCC Code Coverage Report

Directory: src/
File: src/pg/support_pg.c
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 78 80 97.5%
Functions: 19 19 100.0%
Branches: 17 28 60.7%
Line Branch Exec Source
1 /*
2 * Copyright (c) 2026 Tiger Data, Inc.
3 * Licensed under the PostgreSQL License. See LICENSE for details.
4 *
5 * vs_pg.c - pg_vectorsearch PostgreSQL extension entry point
6 */
7
8 #include <postgres.h>
9
10 #include "vs_config.h"
11
12 #include <access/reloptions.h>
13 #include <catalog/namespace.h>
14 #include <fmgr.h>
15 #include <utils/builtins.h>
16 #include <utils/guc.h>
17
18 #include "algo/distance.h"
19 #include "algo/kmeans.h"
20 #include "explain.h"
21 #include "git_commit.h"
22 #include "index/index_build.h"
23 #include "index/posting_page.h"
24 #include "index/query_scan.h"
25 #include "pg/bufstorage.h"
26 #include "scan.h"
27 #include "scan_bound.h"
28 #include "support_pg.h"
29
30 255 PG_MODULE_MAGIC;
31
32 /* GUC variables */
33 int prism_distance_mode = VS_DISTANCE_MODE_DEFAULT;
34 int prism_nprobe = 0;
35 int prism_query_limit = 0;
36 int prism_fastscan_bits = 16;
37 bool prism_rerank = true;
38 bool prism_log_build_stats = false;
39 double prism_centroid_error_scale = 0.0;
40 double prism_centroid_beam_scale = 0.5;
41 int prism_leaf_refine_threshold = 0;
42 static bool prism_recent_buffers = true;
43 double prism_probe_expand = 2.0;
44 static int prism_rerank_pool = 0;
45
46 static void
47 255 vs_recent_buffers_assign_hook(bool newval, void *extra)
48 {
49 255 vs_pg_storage_set_recent_buffers(newval);
50 255 }
51
52 static void
53 257 vs_probe_expand_assign_hook(double newval, void *extra)
54 {
55 257 prism_query_set_probe_expand(newval);
56 257 }
57
58 static void
59 266 vs_rerank_pool_assign_hook(int newval, void *extra)
60 {
61 266 prism_query_set_rerank_pool((int32_t)newval);
62 266 }
63
64 static const struct config_enum_entry vs_distance_mode_options[] = {
65 {"default", VS_DISTANCE_MODE_DEFAULT, false},
66 {"asymmetric", VS_DISTANCE_MODE_ASYMMETRIC, false},
67 {"symmetric", VS_DISTANCE_MODE_SYMMETRIC, false},
68 {NULL, 0, false},
69 };
70
71 static const struct config_enum_entry vs_fastscan_bits_options[] = {
72 {"8", 8, false},
73 {"16", 16, false},
74 {NULL, 0, false},
75 };
76
77 /* Index reloptions */
78 relopt_kind prism_relopt_kind;
79
80 static relopt_enum_elt_def distance_mode_relopt_members[] = {
81 {"asymmetric", VS_DISTANCE_MODE_ASYMMETRIC},
82 {"symmetric", VS_DISTANCE_MODE_SYMMETRIC},
83 {NULL, 0},
84 };
85
86 /* "true"/"false" are accepted aliases for "on"/"off". */
87 static relopt_enum_elt_def centroid_compression_relopt_members[] = {
88 {"auto", PRISM_CENTROID_COMPRESSION_AUTO},
89 {"on", PRISM_CENTROID_COMPRESSION_ON},
90 {"true", PRISM_CENTROID_COMPRESSION_ON},
91 {"off", PRISM_CENTROID_COMPRESSION_OFF},
92 {"false", PRISM_CENTROID_COMPRESSION_OFF},
93 {NULL, 0},
94 };
95
96 /* Shared by the `fastscan` and `centroid_fastscan` options. */
97 static relopt_enum_elt_def fastscan_mode_relopt_members[] = {
98 {"auto", PRISM_FASTSCAN_MODE_AUTO},
99 {"on", PRISM_FASTSCAN_MODE_ON},
100 {"true", PRISM_FASTSCAN_MODE_ON},
101 {"off", PRISM_FASTSCAN_MODE_OFF},
102 {"false", PRISM_FASTSCAN_MODE_OFF},
103 {NULL, 0},
104 };
105
106 void _PG_init(void);
107
108 void
109 _PG_init(void)
110 {
111 255 DefineCustomEnumVariable(
112 VS_GUC_PREFIX ".distance_mode",
113 "RaBitQ distance computation mode.",
114 "default (use index setting), asymmetric, or symmetric",
115 &prism_distance_mode,
116 VS_DISTANCE_MODE_DEFAULT,
117 vs_distance_mode_options,
118 PGC_USERSET,
119 0,
120 NULL,
121 NULL,
122 NULL);
123
124 255 DefineCustomIntVariable(
125 VS_GUC_PREFIX ".nprobe",
126 "Number of clusters to probe per query.",
127 "0 derives it from the index's cluster count "
128 "(~0.5*sqrt(nlist), targeting ~0.95 recall). Lower it for "
129 "speed, raise it for recall.",
130 &prism_nprobe,
131 0,
132 0,
133 10000,
134 PGC_USERSET,
135 0,
136 NULL,
137 NULL,
138 NULL);
139
140 255 DefineCustomIntVariable(
141 VS_GUC_PREFIX ".query_limit",
142 "Caps the top-k an index scan is sized for (0 = no cap).",
143 "A scan sizes its top-k from the LIMIT above it, inflated by "
144 "the planner's selectivity estimate for any filter the "
145 "executor applies above it, and bounded by work_mem -- so a "
146 "query wanting more rows than that budget affords is answered "
147 "by raising work_mem. With no LIMIT to size from, the query "
148 "has asked for every row in order and the scan is sized for "
149 "what work_mem affords or the table's estimated row count, "
150 "whichever is smaller. Set this to cap that: it lowers the "
151 "sizing when it is below what the query asked for and is "
152 "ignored otherwise, which bounds a query that has no LIMIT or "
153 "one set far above the rows actually read. No scan is sized "
154 "below 10 rows.",
155 &prism_query_limit,
156 0,
157 0,
158 INT_MAX,
159 PGC_USERSET,
160 0,
161 NULL,
162 NULL,
163 NULL);
164
165 255 DefineCustomIntVariable(
166 VS_GUC_PREFIX ".leaf_refine_threshold",
167 "Sample-per-leaf count below which a subsampled build refines "
168 "the leaf centroids on the full table.",
169 "The k-means sample trains each leaf's encode reference; the "
170 "reference's error shrinks with the leaf's sample count "
171 "(stderr ~ spread/sqrt(n)), so with enough samples per leaf the "
172 "full-table refine scan buys no recall. Below this many samples "
173 "per leaf the sample mean is noisy and the build re-centers the "
174 "references from the whole table (one extra scan). 0 (the "
175 "default) disables refinement; a large value refines whenever "
176 "the sample was bounded below the table.",
177 &prism_leaf_refine_threshold,
178 0,
179 0,
180 INT_MAX,
181 PGC_USERSET,
182 0,
183 NULL,
184 NULL,
185 NULL);
186
187 255 DefineCustomEnumVariable(
188 VS_GUC_PREFIX ".fastscan_bits",
189 "Fastscan LUT quantization bits.",
190 "8 is faster, 16 is more accurate",
191 &prism_fastscan_bits,
192 16,
193 vs_fastscan_bits_options,
194 PGC_USERSET,
195 0,
196 NULL,
197 NULL,
198 NULL);
199
200 255 DefineCustomBoolVariable(
201 VS_GUC_PREFIX ".rerank",
202 "Enable reranking with exact distances.",
203 NULL,
204 &prism_rerank,
205 true,
206 PGC_USERSET,
207 0,
208 NULL,
209 NULL,
210 NULL);
211
212 255 DefineCustomBoolVariable(
213 VS_GUC_PREFIX ".log_build_stats",
214 "Log per-phase index-build resource statistics.",
215 "When on, each build phase logs its elapsed time, build-heap "
216 "usage, and CPU/maxrss (via the server's ShowUsage), plus a final "
217 "summary, at LOG. Off by default; modeled on the core btree "
218 "log_btree_build_stats developer option. The up-front "
219 "planned-allocation line is logged regardless of this setting.",
220 &prism_log_build_stats,
221 false,
222 PGC_SUSET,
223 GUC_NOT_IN_SAMPLE,
224 NULL,
225 NULL,
226 NULL);
227
228 255 DefineCustomRealVariable(
229 VS_GUC_PREFIX ".centroid_error_scale",
230 "Centroid-search beam width, as a multiple of the RaBitQ "
231 "distance-error margin.",
232 "During centroid routing each candidate centroid has an "
233 "approximate (RaBitQ-quantized) distance plus an error margin; "
234 "this multiplies that margin when deciding which centroids the "
235 "beam keeps at each tree level. 0 (default) ignores the margin "
236 "and "
237 "keeps only the closest centroids by point estimate -- the "
238 "narrowest and fastest beam. 1 widens the beam to also keep "
239 "centroids whose error interval still overlaps the cutoff -- ones "
240 "that might rank among the closest once quantization error is "
241 "accounted for -- the most recall-conservative setting, at the "
242 "cost of scoring more centroids. Larger values widen it further. "
243 "Applies to both compressed centroid formats (RaBitQ and "
244 "FASTSCAN); float/half centroid pages have exact distances (no "
245 "error margin) and are unaffected.",
246 &prism_centroid_error_scale,
247 0.0,
248 0.0,
249 10.0,
250 PGC_USERSET,
251 0,
252 NULL,
253 NULL,
254 NULL);
255
256 255 DefineCustomRealVariable(
257 VS_GUC_PREFIX ".centroid_beam_scale",
258 "Intermediate centroid beam width as a fraction of nprobe.",
259 "Sets the beam width at the intermediate tree levels to this "
260 "fraction of nprobe; the leaf level always returns the full "
261 "nprobe, and the beam never drops below a small floor of "
262 "candidates (or nprobe itself, whichever is less) — at small "
263 "nprobe a scaled-down beam saves next to nothing and "
264 "mis-routes. 0.5 (default) matches the benchmark-tuned "
265 "routing shape. 1.0 keeps the full beam (beam_width = nprobe) "
266 "at every level -- the widest and most recall-conservative "
267 "setting. Smaller values score fewer centroids and are faster "
268 "at high nprobe, with a small recall risk if a near leaf's "
269 "ancestor falls outside the narrowed beam.",
270 &prism_centroid_beam_scale,
271 0.5,
272 0.01,
273 1.0,
274 PGC_USERSET,
275 0,
276 NULL,
277 NULL,
278 NULL);
279
280 255 DefineCustomBoolVariable(
281 VS_GUC_PREFIX ".recent_buffers",
282 "Re-pin index pages via a backend-local buffer-id cache.",
283 "Skips the shared buffer-mapping hash lookup (a large share of "
284 "warm scan CPU) by remembering each block's buffer id and "
285 "re-pinning it with ReadRecentBuffer. Stale ids fall back to a "
286 "normal read and self-heal. Costs 4 bytes per index block per "
287 "backend.",
288 &prism_recent_buffers,
289 true,
290 PGC_USERSET,
291 0,
292 NULL,
293 vs_recent_buffers_assign_hook,
294 NULL);
295
296 255 DefineCustomRealVariable(
297 VS_GUC_PREFIX ".probe_expand",
298 "Probe-candidate expansion factor for exact centroid re-rank.",
299 "Routes ceil(nprobe * expand) leaf candidates through the "
300 "centroid beam, re-ranks them by exact query-centroid distance "
301 "(the full-precision rotated centroid on each cluster's first "
302 "posting page), and scans only the best nprobe in that order. "
303 "Corrects the probe-order noise of compressed (RaBitQ) centroid "
304 "routing; skipped automatically for indexes with exact "
305 "float/half centroids. 1.0 means no expansion (identity). "
306 "Gains saturate around 2 (the default); the extra routed "
307 "candidates are capped at 256, which bounds the overhead at "
308 "large nprobe while retaining nearly all of the recall gain.",
309 &prism_probe_expand,
310 2.0,
311 1.0,
312 16.0,
313 PGC_USERSET,
314 0,
315 NULL,
316 vs_probe_expand_assign_hook,
317 NULL);
318
319 255 DefineCustomIntVariable(
320 VS_GUC_PREFIX ".rerank_pool",
321 "Max candidates to exact-rerank per query (0 = automatic, "
322 "-1 = unlimited).",
323 "Rerank only the most promising candidates by approximate "
324 "distance, bounding the exact-distance heap fetches. 0 (the "
325 "default) caps at 3 * k * nprobe^0.15, growing to 1/8th of "
326 "the candidate buffer when noisy distance estimates flood "
327 "it; -1 reranks every threshold survivor; positive values "
328 "set an absolute cap. The effective cap is never below the "
329 "query's k, so results are never truncated.",
330 &prism_rerank_pool,
331 0,
332 -1,
333 1000000,
334 PGC_USERSET,
335 0,
336 NULL,
337 vs_rerank_pool_assign_hook,
338 NULL);
339
340 255 MarkGUCPrefixReserved(VS_GUC_PREFIX);
341
342 255 vs_cblas_pin_single_thread();
343
344 255 prism_relopt_kind = add_reloption_kind();
345 255 add_enum_reloption(
346 prism_relopt_kind,
347 "distance_mode",
348 "RaBitQ distance computation mode",
349 distance_mode_relopt_members,
350 VS_DISTANCE_MODE_ASYMMETRIC,
351 "symmetric is faster but has larger estimation error",
352 NoLock);
353 255 add_int_reloption(
354 prism_relopt_kind,
355 "fan_out",
356 "Children per tree node (2-255)",
357 PRISM_DEFAULT_FAN_OUT,
358 PRISM_MIN_FAN_OUT,
359 PRISM_MAX_FAN_OUT,
360 NoLock);
361 255 add_int_reloption(
362 prism_relopt_kind,
363 "nlist",
364 "Number of IVF clusters (0 = auto from sqrt(rows))",
365 PRISM_DEFAULT_NLIST,
366 PRISM_MIN_NLIST,
367 PRISM_MAX_NLIST,
368 NoLock);
369 255 add_int_reloption(
370 prism_relopt_kind,
371 "target_pages",
372 "Posting pages each list rests at (1-64)",
373 PRISM_DEFAULT_TARGET_PAGES,
374 PRISM_MIN_TARGET_PAGES,
375 PRISM_MAX_TARGET_PAGES,
376 NoLock);
377 255 add_int_reloption(
378 prism_relopt_kind,
379 "kmeans_nredo",
380 "K-means restarts for cluster quality (1 = no restart)",
381 1,
382 1,
383 20,
384 NoLock);
385 255 add_real_reloption(
386 prism_relopt_kind,
387 "soar_lambda",
388 "SOAR replication lambda (0 = off)",
389 PRISM_DEFAULT_SOAR_LAMBDA,
390 0.0,
391 100.0,
392 NoLock);
393 255 add_real_reloption(
394 prism_relopt_kind,
395 "boundary_epsilon",
396 "Boundary replication gap threshold (0 = off)",
397 PRISM_DEFAULT_BOUNDARY_EPSILON,
398 0.0,
399 100.0,
400 NoLock);
401 255 add_enum_reloption(
402 prism_relopt_kind,
403 "centroid_compression",
404 "RaBitQ compression for centroid pages",
405 centroid_compression_relopt_members,
406 PRISM_CENTROID_COMPRESSION_AUTO,
407 "auto compresses for L2/cosine and skips inner product",
408 NoLock);
409 255 add_enum_reloption(
410 prism_relopt_kind,
411 "fastscan",
412 "VPSHUFB fastscan posting page format",
413 fastscan_mode_relopt_members,
414 PRISM_FASTSCAN_MODE_AUTO,
415 "auto uses it up to the dimension where a group fits a page",
416 NoLock);
417 255 add_enum_reloption(
418 prism_relopt_kind,
419 "centroid_fastscan",
420 "FASTSCAN-format centroid pages",
421 fastscan_mode_relopt_members,
422 PRISM_FASTSCAN_MODE_AUTO,
423 "auto follows the resolved centroid compression and the "
424 "dimension limit; on errors where either is unavailable",
425 NoLock);
426
427 255 vs_distance_init();
428 255 vs_rabitq_init_simd();
429 255 prism_explain_init();
430 255 prism_scan_bound_init();
431 255 }
432
433 /* ----------------------------------------------------------------
434 * Validation helpers
435 * ---------------------------------------------------------------- */
436
437 void
438 46857 vs_pg_check_dim_valid(int dim)
439 {
440
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 46857 times.
 46857 if (dim < 1)
441 ✗ ereport(ERROR,
442 (errcode(ERRCODE_DATA_EXCEPTION),
443 errmsg("vec32 must have at least 1 dimension")));
444
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 46857 times.
 46857 if (dim > VEC32_MAX_DIM)
445 ✗ ereport(ERROR,
446 (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
447 errmsg("vec32 cannot have more than %d dimensions",
448 VEC32_MAX_DIM)));
449 46857 }
450
451 /*
452 * rabitq_params_generate() builds a dim x dim orthogonal matrix by
453 * Gram-Schmidt -- O(dim^3) in time, O(dim^2) in memory -- and is callable by
454 * any role, so an oversized dim is a CPU/memory denial-of-service vector (at
455 * the generic vector cap a single call runs for minutes at 100% CPU). Bound
456 * it by PRISM_INDEX_MAX_DIM: generating params for a dimension no prism index
457 * can hold is pointless, so the largest indexable dimension is the natural
458 * ceiling, and it tracks the index limit automatically.
459 */
460 void
461 9 vs_pg_check_rabitq_params_dim_valid(int dim)
462 {
463 9 vs_pg_check_dim_valid(dim);
464
2/2
✓ Branch 0 taken 1 times.
✓ Branch 1 taken 8 times.
 9 if (dim > PRISM_INDEX_MAX_DIM)
465
1/2
✓ Branch 1 taken 1 times.
✗ Branch 2 not taken.
 1 ereport(ERROR,
466 (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
467 errmsg("cannot generate rabitq_params for more than %d "
468 "dimensions",
469 PRISM_INDEX_MAX_DIM),
470 errhint("Building the transform matrix is O(dim^3); no "
471 "prism index supports more dimensions than this.")));
472 8 }
473
474 void
475 77872 vs_pg_check_dims_match(int dim_a, int dim_b)
476 {
477
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 77869 times.
 77872 if (dim_a != dim_b)
478
1/2
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
 3 ereport(ERROR,
479 (errcode(ERRCODE_DATA_EXCEPTION),
480 errmsg("different vec32 dimensions %d and %d",
481 dim_a,
482 dim_b)));
483 77869 }
484
485 void
486 333327 vs_pg_check_expected_dim(int actual, int expected)
487 {
488
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 333324 times.
 333327 if (expected != -1 && actual != expected)
489
1/2
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
 3 ereport(ERROR,
490 (errcode(ERRCODE_DATA_EXCEPTION),
491 errmsg("expected %d dimensions, not %d", expected, actual)));
492 333324 }
493
494 void
495 18057204 vs_pg_check_value_finite(float val)
496 {
497
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 18057202 times.
 18057204 if (isinf(val))
498
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 ereport(ERROR,
499 (errcode(ERRCODE_DATA_EXCEPTION),
500 errmsg("infinite value not allowed in vec32")));
501
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 18057200 times.
 18057202 if (isnan(val))
502
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 ereport(ERROR,
503 (errcode(ERRCODE_DATA_EXCEPTION),
504 errmsg("NaN value not allowed in vec32")));
505 18057200 }
506
507 /* ----------------------------------------------------------------
508 * Build identity
509 * ---------------------------------------------------------------- */
510
511 /*
512 * Return the git commit the extension was built from. The string is
513 * baked into the binary via vcs_tag at build time (see meson.build).
514 * Rekall and other tooling use this to identify a build for
515 * benchmark reports.
516 */
517 9 PG_FUNCTION_INFO_V1(vs_git_commit);
518
519 Datum
520 2 vs_git_commit(PG_FUNCTION_ARGS)
521 {
522 2 PG_RETURN_TEXT_P(cstring_to_text(VS_GIT_COMMIT));
523 }
524
525 /*
526 * Return the extension version the binary was built as. Comes from
527 * meson.build via vs_config.h, the single source of truth for the
528 * version string.
529 */
530 19 PG_FUNCTION_INFO_V1(vs_extension_version);
531
532 Datum
533 12 vs_extension_version(PG_FUNCTION_ARGS)
534 {
535 12 PG_RETURN_TEXT_P(cstring_to_text(VS_VERSION));
536 }
537
538 /* ----------------------------------------------------------------
539 * Metric identifier support functions (FUNCTION 2 in opclasses)
540 * ---------------------------------------------------------------- */
541
542 45 PG_FUNCTION_INFO_V1(prism_metric_l2);
543
544 Datum
545 180 prism_metric_l2(PG_FUNCTION_ARGS)
546 {
547 180 PG_RETURN_INT32(DISTANCE_L2);
548 }
549
550 9 PG_FUNCTION_INFO_V1(prism_metric_ip);
551
552 Datum
553 3 prism_metric_ip(PG_FUNCTION_ARGS)
554 {
555 3 PG_RETURN_INT32(DISTANCE_INNER_PRODUCT);
556 }
557
558 13 PG_FUNCTION_INFO_V1(prism_metric_cosine);
559
560 Datum
561 13 prism_metric_cosine(PG_FUNCTION_ARGS)
562 {
563 13 PG_RETURN_INT32(DISTANCE_COSINE);
564 }
565