GCC Code Coverage Report

Directory: src/
File: src/pg/rabitq_pg.c
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 148 153 96.7%
Functions: 31 31 100.0%
Branches: 51 82 62.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 * vs_pg_rabitq.c - PostgreSQL functions for the rabitq type
6 *
7 * Type I/O, comparison operators, accessor functions, and encoding.
8 * Text format: {bits:f_add:f_rescale} where bits are 0/1 characters.
9 */
10
11 #include <postgres.h>
12
13 #include <fmgr.h>
14 #include <lib/stringinfo.h>
15 #include <utils/array.h>
16
17 #include "quant/matrix.h"
18 #include "support_pg.h"
19
20 /* ----------------------------------------------------------------
21 * Type I/O
22 * ---------------------------------------------------------------- */
23
24 9 PG_FUNCTION_INFO_V1(vs_rabitq_in);
25
26 Datum
27 75 vs_rabitq_in(PG_FUNCTION_ARGS)
28 {
29 75 char *str = PG_GETARG_CSTRING(0);
30 75 int32 typmod = PG_GETARG_INT32(2);
31 75 uint8_t bits_buf[VS_RABITQ_BYTES(VEC32_MAX_DIM)];
32 75 int dim = 0;
33 75 char *p = str;
34
35 75 memset(bits_buf, 0, sizeof(bits_buf));
36
37 /* Skip whitespace */
38
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 75 times.
 79 while (*p == ' ' || *p == '\t')
39 4 p++;
40
41 /* Expect opening '{' */
42
2/2
✓ Branch 0 taken 1 times.
✓ Branch 1 taken 74 times.
 75 if (*p != '{')
43
1/2
✓ Branch 1 taken 1 times.
✗ Branch 2 not taken.
 1 ereport(ERROR,
44 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
45 errmsg("rabitq must start with \"{\"")));
46 74 p++;
47
48 /* Parse bit characters */
49
4/4
✓ Branch 0 taken 574 times.
✓ Branch 1 taken 72 times.
✓ Branch 2 taken 573 times.
✓ Branch 3 taken 1 times.
 646 while (*p && *p != ':' && *p != '}')
50 {
51
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 573 times.
 573 if (dim >= VEC32_MAX_DIM)
52 ✗ ereport(ERROR,
53 (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
54 errmsg("rabitq cannot have more than %d dimensions",
55 VEC32_MAX_DIM)));
56
57
2/2
✓ Branch 0 taken 261 times.
✓ Branch 1 taken 312 times.
 573 if (*p == '1')
58 261 bits_buf[dim / 8] |= (1 << (dim % 8));
59
2/2
✓ Branch 0 taken 1 times.
✓ Branch 1 taken 311 times.
 312 else if (*p != '0')
60
1/2
✓ Branch 1 taken 1 times.
✗ Branch 2 not taken.
 1 ereport(ERROR,
61 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
62 errmsg("invalid bit character '%c' in rabitq", *p)));
63
64 572 dim++;
65 572 p++;
66 }
67
68
2/2
✓ Branch 0 taken 1 times.
✓ Branch 1 taken 72 times.
 73 if (dim < 1)
69
1/2
✓ Branch 1 taken 1 times.
✗ Branch 2 not taken.
 1 ereport(ERROR,
70 (errcode(ERRCODE_DATA_EXCEPTION),
71 errmsg("rabitq must have at least 1 dimension")));
72
73 /* Expect ':' separator */
74
2/2
✓ Branch 0 taken 1 times.
✓ Branch 1 taken 71 times.
 72 if (*p != ':')
75
1/2
✓ Branch 1 taken 1 times.
✗ Branch 2 not taken.
 1 ereport(ERROR,
76 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
77 errmsg("expected \":\" after bits in rabitq")));
78 71 p++;
79
80 /* Parse f_add */
81 71 char *end;
82 71 errno = 0;
83 71 float f_add = strtof(p, &end);
84
2/4
✓ Branch 0 taken 71 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 71 times.
 71 if (end == p || errno == ERANGE)
85 ✗ ereport(ERROR,
86 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
87 errmsg("invalid f_add value in rabitq: \"%s\"", str)));
88 71 vs_pg_check_value_finite(f_add);
89 70 p = end;
90
91 /* Expect ':' separator */
92
2/2
✓ Branch 0 taken 1 times.
✓ Branch 1 taken 69 times.
 70 if (*p != ':')
93
1/2
✓ Branch 1 taken 1 times.
✗ Branch 2 not taken.
 1 ereport(ERROR,
94 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
95 errmsg("expected \":\" after f_add in rabitq")));
96 69 p++;
97
98 /* Parse f_rescale */
99 69 errno = 0;
100 69 float f_rescale = strtof(p, &end);
101
2/4
✓ Branch 0 taken 69 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 69 times.
 69 if (end == p || errno == ERANGE)
102 ✗ ereport(ERROR,
103 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
104 errmsg("invalid f_rescale value in rabitq: \"%s\"", str)));
105 69 vs_pg_check_value_finite(f_rescale);
106 68 p = end;
107
108 /* Expect closing '}' */
109
2/2
✓ Branch 0 taken 1 times.
✓ Branch 1 taken 67 times.
 68 if (*p != '}')
110
1/2
✓ Branch 1 taken 1 times.
✗ Branch 2 not taken.
 1 ereport(ERROR,
111 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
112 errmsg("rabitq must end with \"}\"")));
113 67 p++;
114
115 /* Check no trailing content */
116
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 67 times.
 71 while (*p == ' ' || *p == '\t')
117 4 p++;
118
2/2
✓ Branch 0 taken 1 times.
✓ Branch 1 taken 66 times.
 67 if (*p != '\0')
119
1/2
✓ Branch 1 taken 1 times.
✗ Branch 2 not taken.
 1 ereport(ERROR,
120 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
121 errmsg("unexpected characters after \"}\" in "
122 "rabitq")));
123
124 66 vs_pg_check_expected_dim(dim, typmod);
125
126 66 RaBitQVector *result = vs_pg_rabitq_alloc(dim);
127 66 result->f_add = f_add;
128 66 result->f_rescale = f_rescale;
129 66 memcpy(result->bits, bits_buf, VS_RABITQ_BYTES(dim));
130
131 66 PG_RETURN_RABITQ_P(result);
132 }
133
134 9 PG_FUNCTION_INFO_V1(vs_rabitq_out);
135
136 Datum
137 16 vs_rabitq_out(PG_FUNCTION_ARGS)
138 {
139 16 RaBitQVector *v = PG_GETARG_RABITQ_P(0);
140 16 StringInfoData buf;
141
142 16 initStringInfo(&buf);
143 16 appendStringInfoChar(&buf, '{');
144
145
2/2
✓ Branch 1 taken 121 times.
✓ Branch 2 taken 16 times.
 137 for (int i = 0; i < v->dim; i++)
146 121 appendStringInfoChar(
147
2/2
✓ Branch 0 taken 60 times.
✓ Branch 1 taken 61 times.
 121 &buf, (v->bits[i / 8] >> (i % 8)) & 1 ? '1' : '0');
148
149 16 appendStringInfo(&buf, ":%.8g:%.8g}", v->f_add, v->f_rescale);
150
151 16 PG_RETURN_CSTRING(buf.data);
152 }
153
154 9 PG_FUNCTION_INFO_V1(vs_rabitq_typmod_in);
155
156 Datum
157 4 vs_rabitq_typmod_in(PG_FUNCTION_ARGS)
158 {
159 4 ArrayType *ta = PG_GETARG_ARRAYTYPE_P(0);
160 4 int n;
161 4 int32 *tl = ArrayGetIntegerTypmods(ta, &n);
162
163
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (n != 1)
164 ✗ ereport(ERROR,
165 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
166 errmsg("invalid type modifier")));
167
168 4 int dim = tl[0];
169
170 4 vs_pg_check_dim_valid(dim);
171
172 4 PG_RETURN_INT32(dim);
173 }
174
175 /* ----------------------------------------------------------------
176 * Accessor functions
177 * ---------------------------------------------------------------- */
178
179 9 PG_FUNCTION_INFO_V1(vs_rabitq_dims);
180
181 Datum
182 4 vs_rabitq_dims(PG_FUNCTION_ARGS)
183 {
184 4 RaBitQVector *v = PG_GETARG_RABITQ_P(0);
185 4 PG_RETURN_INT32(v->dim);
186 }
187
188 9 PG_FUNCTION_INFO_V1(vs_rabitq_f_add);
189
190 Datum
191 2 vs_rabitq_f_add(PG_FUNCTION_ARGS)
192 {
193 2 RaBitQVector *v = PG_GETARG_RABITQ_P(0);
194 2 PG_RETURN_FLOAT8((double)v->f_add);
195 }
196
197 9 PG_FUNCTION_INFO_V1(vs_rabitq_f_rescale);
198
199 Datum
200 2 vs_rabitq_f_rescale(PG_FUNCTION_ARGS)
201 {
202 2 RaBitQVector *v = PG_GETARG_RABITQ_P(0);
203 2 PG_RETURN_FLOAT8((double)v->f_rescale);
204 }
205
206 /* ----------------------------------------------------------------
207 * Comparison functions (for btree opclass)
208 *
209 * Compare dim first, then memcmp on bits. f_add/f_rescale are
210 * intentionally excluded -- bits are the quantized identity.
211 * ---------------------------------------------------------------- */
212
213 static int
214 17 vs_rabitq_cmp_internal(RaBitQVector *a, RaBitQVector *b)
215 {
216
2/2
✓ Branch 0 taken 1 times.
✓ Branch 1 taken 16 times.
 17 if (a->dim != b->dim)
217
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 times.
 1 return (a->dim < b->dim) ? -1 : 1;
218
219 16 int nbytes = VS_RABITQ_BYTES(a->dim);
220 16 return memcmp(a->bits, b->bits, nbytes);
221 }
222
223 9 PG_FUNCTION_INFO_V1(vs_rabitq_cmp);
224
225 Datum
226 9 vs_rabitq_cmp(PG_FUNCTION_ARGS)
227 {
228 9 RaBitQVector *a = PG_GETARG_RABITQ_P(0);
229 9 RaBitQVector *b = PG_GETARG_RABITQ_P(1);
230 9 PG_RETURN_INT32(vs_rabitq_cmp_internal(a, b));
231 }
232
233 9 PG_FUNCTION_INFO_V1(vs_rabitq_lt);
234
235 Datum
236 2 vs_rabitq_lt(PG_FUNCTION_ARGS)
237 {
238 2 RaBitQVector *a = PG_GETARG_RABITQ_P(0);
239 2 RaBitQVector *b = PG_GETARG_RABITQ_P(1);
240
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 PG_RETURN_BOOL(vs_rabitq_cmp_internal(a, b) < 0);
241 }
242
243 9 PG_FUNCTION_INFO_V1(vs_rabitq_le);
244
245 Datum
246 1 vs_rabitq_le(PG_FUNCTION_ARGS)
247 {
248 1 RaBitQVector *a = PG_GETARG_RABITQ_P(0);
249 1 RaBitQVector *b = PG_GETARG_RABITQ_P(1);
250
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 times.
 1 PG_RETURN_BOOL(vs_rabitq_cmp_internal(a, b) <= 0);
251 }
252
253 9 PG_FUNCTION_INFO_V1(vs_rabitq_eq);
254
255 Datum
256 2 vs_rabitq_eq(PG_FUNCTION_ARGS)
257 {
258 2 RaBitQVector *a = PG_GETARG_RABITQ_P(0);
259 2 RaBitQVector *b = PG_GETARG_RABITQ_P(1);
260
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 PG_RETURN_BOOL(vs_rabitq_cmp_internal(a, b) == 0);
261 }
262
263 9 PG_FUNCTION_INFO_V1(vs_rabitq_ne);
264
265 Datum
266 1 vs_rabitq_ne(PG_FUNCTION_ARGS)
267 {
268 1 RaBitQVector *a = PG_GETARG_RABITQ_P(0);
269 1 RaBitQVector *b = PG_GETARG_RABITQ_P(1);
270
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 times.
 1 PG_RETURN_BOOL(vs_rabitq_cmp_internal(a, b) != 0);
271 }
272
273 9 PG_FUNCTION_INFO_V1(vs_rabitq_ge);
274
275 Datum
276 1 vs_rabitq_ge(PG_FUNCTION_ARGS)
277 {
278 1 RaBitQVector *a = PG_GETARG_RABITQ_P(0);
279 1 RaBitQVector *b = PG_GETARG_RABITQ_P(1);
280
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 times.
 1 PG_RETURN_BOOL(vs_rabitq_cmp_internal(a, b) >= 0);
281 }
282
283 9 PG_FUNCTION_INFO_V1(vs_rabitq_gt);
284
285 Datum
286 1 vs_rabitq_gt(PG_FUNCTION_ARGS)
287 {
288 1 RaBitQVector *a = PG_GETARG_RABITQ_P(0);
289 1 RaBitQVector *b = PG_GETARG_RABITQ_P(1);
290
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 times.
 1 PG_RETURN_BOOL(vs_rabitq_cmp_internal(a, b) > 0);
291 }
292
293 /* ----------------------------------------------------------------
294 * Cast function (typmod enforcement)
295 * ---------------------------------------------------------------- */
296
297 9 PG_FUNCTION_INFO_V1(vs_rabitq);
298
299 Datum
300 6 vs_rabitq(PG_FUNCTION_ARGS)
301 {
302 6 RaBitQVector *v = PG_GETARG_RABITQ_P(0);
303 6 int32 typmod = PG_GETARG_INT32(1);
304
305 6 vs_pg_check_expected_dim(v->dim, typmod);
306
307 5 PG_RETURN_RABITQ_P(v);
308 }
309
310 /* ----------------------------------------------------------------
311 * Encode function
312 * ---------------------------------------------------------------- */
313
314 9 PG_FUNCTION_INFO_V1(vs_rabitq_encode_pg);
315
316 Datum
317 5 vs_rabitq_encode_pg(PG_FUNCTION_ARGS)
318 {
319 5 Vec32 *input = PG_GETARG_VEC32_P(0);
320 5 Vec32 *centroid = PG_GETARG_VEC32_P(1);
321 5 RaBitQParamsPG *params = PG_GETARG_RABITQ_PARAMS_P(2);
322
323 5 vs_pg_check_dims_match(input->dim, centroid->dim);
324 4 vs_pg_check_dims_match(input->dim, params->dim);
325
326 4 int dim = input->dim;
327
328 /* Build RaBitQParams from the PG varlena's pre-computed matrix */
329 4 RaBitQParams *rparams =
330 4 vs_rabitq_create_from_matrix(dim, params->seed, params->P);
331
332 4 RaBitQVector *result = vs_pg_rabitq_alloc(dim);
333
334 4 Vec32Ref input_ref = Vec32ToRef(input);
335 4 Vec32Ref centroid_ref = Vec32ToRef(centroid);
336 4 RaBitQData *data = VS_RABITQ_DATA(result);
337
338 4 int ret = vs_rabitq_encode_into(rparams, input_ref, centroid_ref, data);
339
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (ret != 0)
340 ✗ ereport(ERROR,
341 (errcode(ERRCODE_INTERNAL_ERROR),
342 errmsg("rabitq encoding failed")));
343
344 4 PG_RETURN_RABITQ_P(result);
345 }
346