GCC Code Coverage Report

Directory: src/
File: src/standalone/api.c
Date: 2026-10-23 09:17:32
Exec Total Coverage
Lines: 70 86 81.4%
Functions: 8 8 100.0%
Branches: 30 46 65.2%
Line Branch Exec Source
1 /*
2 * Copyright (c) 2026 Tiger Data, Inc.
3 * Licensed under the PostgreSQL License. See LICENSE for details.
4 *
5 * api.c - Public C API for standalone pg_vectorsearch library
6 *
7 * Thin wrapper around PrismIndex (standalone/index.h) and PrismQueryCtx
8 * (standalone/query.h). The handle bundles both so callers get a
9 * single opaque pointer.
10 *
11 * A top-level memory context owns the handle. The index and query
12 * context create their own child contexts internally.
13 */
14
15 #include <string.h>
16
17 #include "core/memory.h"
18 #include "index/query_scan.h"
19 #include "standalone/api.h"
20 #include "standalone/index.h"
21 #include "standalone/query.h"
22
23 struct VsHandle
24 {
25 VsMemCtx memctx; /* top-level context, owns everything */
26 PrismIndex *idx;
27 PrismQueryCtx *qctx;
28 };
29
30 /* ----------------------------------------------------------------
31 * Parse helpers
32 * ---------------------------------------------------------------- */
33
34 static DistanceMetric
35 14 parse_metric(const char *s)
36 {
37
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 if (s == NULL)
38 ✗ return DISTANCE_L2;
39
3/4
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
 14 if (strcmp(s, "angular") == 0 || strcmp(s, "cosine") == 0)
40 2 return DISTANCE_COSINE;
41 12 return DISTANCE_L2;
42 }
43
44 static PrismCentroidFormat
45 14 parse_centroid_fmt(const char *s)
46 {
47
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 if (s == NULL)
48 ✗ return PRISM_CENTROID_FMT_RABITQ;
49
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 12 times.
 14 if (strcmp(s, "float32") == 0)
50 2 return PRISM_CENTROID_FMT_FLOAT;
51
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
 12 if (strcmp(s, "float16") == 0)
52 ✗ return PRISM_CENTROID_FMT_HALF;
53
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
 12 if (strcmp(s, "fastscan") == 0)
54 ✗ return PRISM_CENTROID_FMT_FASTSCAN;
55 12 return PRISM_CENTROID_FMT_RABITQ;
56 }
57
58 static PrismPostingFormat
59 14 parse_posting_fmt(const char *s)
60 {
61
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 14 if (s != NULL && strcmp(s, "flat") == 0)
62 ✗ return PRISM_POSTING_FMT_FLAT;
63 14 return PRISM_POSTING_FMT_PAGES; /* default: pages (matches PG on-disk) */
64 }
65
66 static VsDistanceMode
67 10 parse_distance_mode(const char *s)
68 {
69
3/4
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 8 times.
 10 if (s != NULL && strcmp(s, "symmetric") == 0)
70 2 return VS_DISTANCE_MODE_SYMMETRIC;
71 8 return VS_DISTANCE_MODE_ASYMMETRIC;
72 }
73
74 /* ----------------------------------------------------------------
75 * API
76 * ---------------------------------------------------------------- */
77
78 VsHandle *
79 14 vs_handle_create(
80 Vec32Source *src,
81 uint32_t nlist,
82 uint32_t fan_out,
83 const char *metric,
84 const char *centroid_fmt,
85 const char *posting_fmt,
86 uint32_t km_nredo,
87 uint32_t km_max_iter,
88 double soar_lambda,
89 double boundary_epsilon,
90 int fastscan,
91 int32_t nworkers,
92 uint32_t max_k,
93 uint32_t max_nprobe,
94 PrismBuildInfo *info)
95 {
96 /* Create context as child of the current context if one exists,
97 * otherwise as a top-level context. This works both when called
98 * from the CLI (which sets up a cli_memctx) and from Python
99 * ctypes (where no context is set). */
100 14 VsMemCtx parent = vs_current_memctx;
101 14 VsMemCtx memctx = vs_memctx_create(parent, "vs_handle");
102 14 VsMemCtx old_ctx = vs_memctx_switch(memctx);
103
104 14 PrismCentroidFormat fmt = parse_centroid_fmt(centroid_fmt);
105
106 42 PrismIndexConfig config = {
107 .nlist = nlist,
108 .fan_out = fan_out,
109 .centroid_fmt = fmt,
110 14 .metric = parse_metric(metric),
111 .km_nredo = km_nredo,
112 .km_max_iter = km_max_iter,
113 .soar_lambda = soar_lambda,
114 .boundary_epsilon = boundary_epsilon,
115 .fastscan = fastscan,
116 .nworkers = nworkers,
117 .encode_rabitq = true,
118 14 .posting_fmt = parse_posting_fmt(posting_fmt),
119 };
120
121 PrismBuildStats build_stats;
122 14 PrismIndex *idx = prism_index_build(src, &config, &build_stats);
123
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 if (idx == NULL)
124 {
125 ✗ vs_memctx_switch(old_ctx);
126 ✗ vs_memctx_delete(memctx);
127 ✗ return NULL;
128 }
129
130 /* Fill build stats if requested */
131
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 10 times.
 14 if (info != NULL)
132 {
133 4 info->stats = build_stats;
134 4 info->nlist = idx->nlist;
135 4 info->nlevels = idx->base.nlevels;
136 4 info->nvecs = idx->nvecs;
137 4 info->max_cluster = 0;
138 4 info->min_cluster = UINT32_MAX;
139
2/2
✓ Branch 0 taken 36 times.
✓ Branch 1 taken 4 times.
 40 for (uint32_t c = 0; c < idx->nlist; c++)
140 {
141 36 uint32_t sz = idx->clusters[c].count;
142
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 36 times.
 36 if (sz > info->max_cluster)
143 ✗ info->max_cluster = sz;
144
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 32 times.
 36 if (sz < info->min_cluster)
145 4 info->min_cluster = sz;
146 }
147
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (idx->nlist == 0)
148 {
149 /* No centroids, so there is nothing to estimate from. */
150 ✗ info->min_cluster = 0;
151 }
152
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 else if (info->max_cluster == 0)
153 {
154 /* Pages mode skips cluster list construction —
155 * fall back to estimates */
156 4 info->max_cluster = (idx->nvecs + idx->nlist - 1) / idx->nlist;
157 4 info->min_cluster = idx->nvecs / idx->nlist;
158 }
159 }
160
161 /* 0 nprobe is the same auto rule the bench and the extension use.
162 * The context is sized once, to the query this handle will run. */
163
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 if (max_k == 0)
164 ✗ max_k = 1;
165
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 if (max_nprobe == 0)
166 ✗ max_nprobe = prism_auto_nprobe(idx->nlist);
167
168 14 PrismQueryCtx *qctx = prism_query_ctx_create(idx, max_k, max_nprobe);
169
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 if (qctx == NULL)
170 {
171 ✗ prism_index_destroy(idx);
172 ✗ vs_memctx_switch(old_ctx);
173 ✗ vs_memctx_delete(memctx);
174 ✗ return NULL;
175 }
176
177 14 VsHandle *handle = vs_alloc(sizeof(VsHandle));
178 14 handle->memctx = memctx;
179 14 handle->idx = idx;
180 14 handle->qctx = qctx;
181
182 14 vs_memctx_switch(old_ctx);
183
184 14 return handle;
185 }
186
187 VsHandle *
188 14 vs_handle_create_from_array(
189 const float *vectors,
190 uint32_t nvecs,
191 uint32_t dim,
192 uint32_t nlist,
193 uint32_t fan_out,
194 const char *metric,
195 const char *centroid_fmt,
196 const char *posting_fmt,
197 uint32_t km_nredo,
198 uint32_t km_max_iter,
199 double soar_lambda,
200 double boundary_epsilon,
201 int fastscan,
202 int32_t nworkers,
203 uint32_t max_k,
204 uint32_t max_nprobe,
205 PrismBuildInfo *info)
206 {
207 VsArraySource array_src;
208 14 vs_array_source_init(&array_src, vectors, nvecs, dim);
209 14 return vs_handle_create(
210 &array_src.base,
211 nlist,
212 fan_out,
213 metric,
214 centroid_fmt,
215 posting_fmt,
216 km_nredo,
217 km_max_iter,
218 soar_lambda,
219 boundary_epsilon,
220 fastscan,
221 nworkers,
222 max_k,
223 max_nprobe,
224 info);
225 }
226
227 uint32_t
228 12 vs_handle_query(
229 VsHandle *handle,
230 const float *query,
231 uint32_t k,
232 uint32_t nprobe,
233 const char *distance_mode,
234 bool rerank,
235 uint32_t *result_ids)
236 {
237
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 10 times.
 12 if (handle == NULL)
238 2 return 0;
239
240 10 VsDistanceMode mode = parse_distance_mode(distance_mode);
241
242 10 return prism_query_exec(
243 handle->qctx, query, k, nprobe, mode, rerank, result_ids);
244 }
245
246 void
247 16 vs_handle_destroy(VsHandle *handle)
248 {
249
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 14 times.
 16 if (handle == NULL)
250 2 return;
251
252 /* Destroy query ctx and index first (they manage their own
253 * child contexts), then delete the top-level context which
254 * frees the handle itself. */
255 14 prism_query_ctx_destroy(handle->qctx);
256 14 prism_index_destroy(handle->idx);
257
258 14 VsMemCtx memctx = handle->memctx;
259 14 vs_memctx_delete(memctx);
260 }
261