GCC Code Coverage Report

Directory: src/
File: src/index/centroid_build.c
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 103 114 90.4%
Functions: 2 2 100.0%
Branches: 40 56 71.4%
Line Branch Exec Source
1 /*
2 * Copyright (c) 2026 Tiger Data, Inc.
3 * Licensed under the PostgreSQL License. See LICENSE for details.
4 *
5 * centroid_build.c - Generic centroid page writer
6 *
7 * Writes centroid entries to linked pages via VsStorage. Format-
8 * agnostic: page format determines metadata and data sizes.
9 */
10
11 #include <assert.h>
12 #include <math.h>
13 #include <string.h>
14
15 #include "core/memory.h"
16 #include "index/centroid_build.h"
17 #include "quant/fastscan.h"
18 #include "quant/rabitq.h"
19
20 BlockNumber
21 710 prism_centroid_write_pages(
22 VsStorage *storage,
23 Dimension dim,
24 uint32_t nlist,
25 PrismCentroidFormat fmt,
26 uint8_t level,
27 uint16_t flags,
28 uint16_t child_count,
29 CentroidEncoder *encoder,
30 const BlockNumber *child_blknos,
31 const float *pt_centroids,
32 BlockNumber start_blkno)
33 {
34 710 bool reserved = (start_blkno != InvalidBlockNumber);
35
1/2
✓ Branch 0 taken 558 times.
✗ Branch 1 not taken.
 710 BlockNumber first_blkno = reserved ? start_blkno : InvalidBlockNumber;
36 710 BlockNumber prev_blkno = InvalidBlockNumber;
37 710 BlockNumber next_blkno = start_blkno;
38 710 Page cur_page = NULL;
39 710 BlockNumber cur_blkno = InvalidBlockNumber;
40
41 152 (void)pt_centroids; /* pt_centroids stored on posting pages, not here */
42
43
2/2
✓ Branch 0 taken 4357 times.
✓ Branch 1 taken 710 times.
 5067 for (uint32_t i = 0; i < nlist; i++)
44 {
45 /* Allocate a new page if needed */
46
4/4
✓ Branch 0 taken 3647 times.
✓ Branch 1 taken 710 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 1067 times.
 5426 if (cur_page == NULL ||
47
2/2
✓ Branch 1 taken 24 times.
✓ Branch 2 taken 2554 times.
 3647 !prism_centroid_page_has_room(cur_page, dim, false))
48 {
49 /* Commit the previous page if any */
50
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 558 times.
 584 if (cur_page != NULL)
51 26 vs_storage_commit_page(storage, cur_blkno);
52
53
1/2
✓ Branch 0 taken 736 times.
✗ Branch 1 not taken.
 736 if (reserved)
54 {
55 736 cur_blkno = next_blkno++;
56 736 cur_page = vs_storage_write_page(storage, cur_blkno);
57 }
58 else
59 {
60 ✗ cur_page = vs_storage_new_page(storage, &cur_blkno);
61 }
62 736 prism_centroid_page_init_fmt(cur_page, level, fmt);
63
64
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 736 times.
 736 if (first_blkno == InvalidBlockNumber)
65 ✗ first_blkno = cur_blkno;
66
67 /* Link previous page to this one. The storage layer
68 * holds at most one buffer, so we commit the new page
69 * first, reopen the previous page to set next_blkno,
70 * then reopen the new page. */
71
2/2
✓ Branch 0 taken 26 times.
✓ Branch 1 taken 710 times.
 736 if (prev_blkno != InvalidBlockNumber)
72 {
73 26 vs_storage_commit_page(storage, cur_blkno);
74
75 26 Page prev_page = vs_storage_write_page(storage, prev_blkno);
76 26 PRISM_CENTROID_OPAQUE(prev_page)->next_blkno = cur_blkno;
77 26 vs_storage_commit_page(storage, prev_blkno);
78
79 26 cur_page = vs_storage_write_page(storage, cur_blkno);
80 }
81
82 736 prev_blkno = cur_blkno;
83 }
84
85 9935 BlockNumber entry_child = child_blknos != NULL ? child_blknos[i]
86
1/2
✓ Branch 0 taken 4357 times.
✗ Branch 1 not taken.
 4357 : InvalidBlockNumber;
87
88 4357 void *dest = prism_centroid_page_add_entry_begin(
89 cur_page, dim, entry_child, child_count, flags);
90
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4357 times.
 4357 assert(dest != NULL);
91 4357 encoder->ops->encode_into(encoder, i, dest);
92 }
93
94 /* Commit last page */
95
1/2
✓ Branch 0 taken 710 times.
✗ Branch 1 not taken.
 710 if (cur_page != NULL)
96 710 vs_storage_commit_page(storage, cur_blkno);
97
98 710 return first_blkno;
99 }
100
101 /* ----------------------------------------------------------------
102 * FASTSCAN centroid writer
103 *
104 * Walks the input centroid list in 32-vector groups. For each group:
105 * 1. Encode each centroid with vs_rabitq_encode_into to get
106 * f_add, f_rescale, and 1-bit packed code bytes.
107 * 2. Derive f_error from (f_add, f_rescale).
108 * 3. Repack the bits into the kPerm0 layout that
109 * vs_fastscan_accumulate_hacc expects, and place the per-
110 * group f_add / f_rescale / f_error / child_blkno arrays in
111 * the group section.
112 *
113 * One page holds `groups_per_page` groups; once a page is full we
114 * commit it and chain via next_blkno, the same way the RABITQ
115 * writer above does. Partial trailing group is padded with
116 * InvalidBlockNumber and zero codes.
117 * ---------------------------------------------------------------- */
118
119 BlockNumber
120 1402 prism_centroid_write_fastscan_pages(
121 VsStorage *storage,
122 Dimension dim,
123 uint32_t nlist,
124 uint8_t level,
125 uint16_t flags,
126 const RaBitQParams *params,
127 const float *vectors,
128 const float *global_mean,
129 const BlockNumber *child_blknos,
130 BlockNumber start_blkno)
131 {
132 1294 (void)flags; /* per-entry flags not stored in FASTSCAN format */
133
134 1402 uint32_t packed_bytes = (dim + 7) / 8;
135
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1294 times.
 1402 uint32_t groups_per_page = prism_centroid_fastscan_max_groups(dim);
136
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1402 times.
 1402 if (groups_per_page == 0)
137 ✗ groups_per_page = 1; /* defensive — bigger dims may need split */
138
139 1402 uint32_t ngroups = (nlist + VS_FASTSCAN_GROUP - 1) / VS_FASTSCAN_GROUP;
140
141 /* Scratch for one group's RaBitQ data (32 entries). */
142 1402 size_t rdata_size = VS_RABITQ_DATA_SIZE(dim);
143 1402 uint8_t *rdata_buf = vs_alloc((size_t)VS_FASTSCAN_GROUP * rdata_size);
144 1402 uint8_t *bits_buf = vs_alloc((size_t)VS_FASTSCAN_GROUP * packed_bytes);
145
146 1402 Vec32Ref mref = {.data = global_mean, .dim = dim};
147
148 1402 bool reserved = (start_blkno != InvalidBlockNumber);
149
1/2
✓ Branch 0 taken 108 times.
✗ Branch 1 not taken.
 1402 BlockNumber first_blkno = reserved ? start_blkno : InvalidBlockNumber;
150 1402 BlockNumber prev_blkno = InvalidBlockNumber;
151 1402 BlockNumber next_blkno = start_blkno;
152 1402 Page cur_page = NULL;
153 1402 BlockNumber cur_blkno = InvalidBlockNumber;
154 1402 uint32_t cur_ngroups = 0;
155
156
2/2
✓ Branch 0 taken 1426 times.
✓ Branch 1 taken 1402 times.
 2828 for (uint32_t g = 0; g < ngroups; g++)
157 {
158 /* Open a fresh page when needed */
159
3/4
✓ Branch 0 taken 1318 times.
✓ Branch 1 taken 108 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 24 times.
 1426 if (cur_page == NULL || cur_ngroups == groups_per_page)
160 {
161
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1402 times.
 1402 if (cur_page != NULL)
162 ✗ vs_storage_commit_page(storage, cur_blkno);
163
164
1/2
✓ Branch 0 taken 1402 times.
✗ Branch 1 not taken.
 1402 if (reserved)
165 {
166 1402 cur_blkno = next_blkno++;
167 1402 cur_page = vs_storage_write_page(storage, cur_blkno);
168 }
169 else
170 ✗ cur_page = vs_storage_new_page(storage, &cur_blkno);
171
172 1402 prism_centroid_page_init_fmt(
173 cur_page, level, PRISM_CENTROID_FMT_FASTSCAN);
174
175
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1402 times.
 1402 if (first_blkno == InvalidBlockNumber)
176 ✗ first_blkno = cur_blkno;
177
178
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1402 times.
 1402 if (prev_blkno != InvalidBlockNumber)
179 {
180 ✗ vs_storage_commit_page(storage, cur_blkno);
181 ✗ Page prev = vs_storage_write_page(storage, prev_blkno);
182 ✗ PRISM_CENTROID_OPAQUE(prev)->next_blkno = cur_blkno;
183 ✗ vs_storage_commit_page(storage, prev_blkno);
184 ✗ cur_page = vs_storage_write_page(storage, cur_blkno);
185 }
186
187 1402 prev_blkno = cur_blkno;
188 1402 cur_ngroups = 0;
189 }
190
191 /* Encode this group's 32 entries (or fewer for the last). */
192 1426 uint32_t g_start = g * VS_FASTSCAN_GROUP;
193 1426 uint32_t g_count = nlist - g_start;
194
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 108 times.
 1426 if (g_count > VS_FASTSCAN_GROUP)
195 24 g_count = VS_FASTSCAN_GROUP;
196
197 /* Get per-entry RaBitQ encoding into rdata_buf, copy bits
198 * out into a flat packed_bytes-stride array for the packer. */
199
2/2
✓ Branch 0 taken 10916 times.
✓ Branch 1 taken 1426 times.
 12342 for (uint32_t v = 0; v < g_count; v++)
200 {
201 10916 Vec32Ref vref = {
202 10916 .data = vectors + (size_t)(g_start + v) * dim,
203 .dim = dim,
204 };
205 10916 RaBitQData *d = (RaBitQData *)(rdata_buf + (size_t)v * rdata_size);
206 10916 vs_rabitq_encode_into(params, vref, mref, d);
207 10916 memcpy(bits_buf + (size_t)v * packed_bytes, d->bits, packed_bytes);
208 }
209
210 /* Write per-entry scalars + child_blkno arrays into the
211 * group section, then pack codes. */
212 1426 uint32_t cur_group_idx = cur_ngroups;
213 1426 char *content = (char *)PageGetContents(cur_page);
214
215 1426 BlockNumber *child = prism_centroid_fastscan_group_child(
216 content, cur_group_idx, dim);
217 1426 float *f_add_arr = prism_centroid_fastscan_group_f_add(
218 content, cur_group_idx, dim);
219 1426 float *f_rescale_arr = prism_centroid_fastscan_group_f_rescale(
220 content, cur_group_idx, dim);
221 1426 float *f_error_arr = prism_centroid_fastscan_group_f_error(
222 content, cur_group_idx, dim);
223 1426 uint8_t *codes = prism_centroid_fastscan_group_codes(
224 content, cur_group_idx, dim);
225
226
2/2
✓ Branch 0 taken 45632 times.
✓ Branch 1 taken 1426 times.
 47058 for (uint32_t v = 0; v < VS_FASTSCAN_GROUP; v++)
227 {
228
2/2
✓ Branch 0 taken 10916 times.
✓ Branch 1 taken 34716 times.
 45632 if (v < g_count)
229 {
230 10916 const RaBitQData *d = (const RaBitQData *)(rdata_buf +
231 10916 (size_t)v *
232 rdata_size);
233 31414 child[v] = child_blknos != NULL ? child_blknos[g_start + v]
234
1/2
✓ Branch 0 taken 10916 times.
✗ Branch 1 not taken.
 10916 : InvalidBlockNumber;
235 10916 f_add_arr[v] = d->f_add;
236 10916 f_rescale_arr[v] = d->f_rescale;
237 10916 f_error_arr[v] =
238 10916 vs_rabitq_derive_f_error(d->f_add, d->f_rescale, dim);
239 }
240 else
241 {
242 34716 child[v] = InvalidBlockNumber;
243 34716 f_add_arr[v] = 0.0f;
244 34716 f_rescale_arr[v] = 0.0f;
245 34716 f_error_arr[v] = 0.0f;
246 }
247 }
248
249 1426 vs_fastscan_pack_codes(bits_buf, g_count, dim, codes);
250
251 1426 PrismCentroidPageOpaque *op = PRISM_CENTROID_OPAQUE(cur_page);
252 1426 op->entry_count += (uint16_t)g_count;
253 1426 cur_ngroups++;
254
255 /* Bump pd_lower so PostgreSQL's hole-compression preserves
256 * our writes. FASTSCAN doesn't use the backward data region
257 * — everything lives in the forward content area. */
258 1426 PageHeader header = (PageHeader)cur_page;
259 1426 header->pd_lower += prism_centroid_fastscan_group_bytes(dim);
260 }
261
262
1/2
✓ Branch 0 taken 1402 times.
✗ Branch 1 not taken.
 1402 if (cur_page != NULL)
263 1402 vs_storage_commit_page(storage, cur_blkno);
264
265 1402 vs_free(rdata_buf);
266 1402 vs_free(bits_buf);
267
268 1402 return first_blkno;
269 }
270