GCC Code Coverage Report

Directory: src/
File: src/index/centroid_build.h
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 30 33 90.9%
Functions: 4 4 100.0%
Branches: 3 5 60.0%
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.h - Generic centroid page writer
6 *
7 * Writes centroid entries to linked pages via VsStorage. Supports
8 * all centroid formats (RaBitQ, float32, float16). Standalone-
9 * compatible — all dependencies are portable.
10 */
11
12 #ifndef PRISM_CENTROID_BUILD_H
13 #define PRISM_CENTROID_BUILD_H
14
15 #include "index/centroid_page.h"
16 #include "index/storage.h"
17
18 /* ----------------------------------------------------------------
19 * CentroidEncoder — vtable for encoding entries into page memory
20 *
21 * Format-specific structs embed CentroidEncoder as their first
22 * field, adding context (vectors, params). write_pages calls
23 * encode_into() for each entry to write directly into a
24 * pre-reserved page data region — no intermediate allocation
25 * or copy needed.
26 * ---------------------------------------------------------------- */
27 typedef struct CentroidEncoder CentroidEncoder;
28
29 typedef struct CentroidEncoderOps
30 {
31 void (*encode_into)(CentroidEncoder *enc, uint32_t index, void *dest);
32 } CentroidEncoderOps;
33
34 struct CentroidEncoder
35 {
36 const CentroidEncoderOps *ops;
37 };
38
39 /* ----------------------------------------------------------------
40 * Built-in encoders for float32, float16, and RaBitQ formats
41 *
42 * Float: memcpy from input array into dest
43 * Half: float→half conversion directly into dest
44 * RaBitQ: vs_rabitq_encode_into directly into dest
45 * ---------------------------------------------------------------- */
46
47 /* Float encoder */
48 typedef struct
49 {
50 CentroidEncoder base;
51 const float *vectors;
52 Dimension dim;
53 } CentroidEncoderFloat;
54
55 static inline void
56 1809 centroid_encode_float(CentroidEncoder *enc, uint32_t index, void *dest)
57 {
58 1809 CentroidEncoderFloat *fe = (CentroidEncoderFloat *)enc;
59 1809 memcpy(dest,
60 1809 fe->vectors + (size_t)index * fe->dim,
61 1809 fe->dim * sizeof(float));
62 1809 }
63
64 static const CentroidEncoderOps centroid_encoder_float_ops = {
65 .encode_into = centroid_encode_float,
66 };
67
68 /* Half encoder */
69 typedef struct
70 {
71 CentroidEncoder base;
72 const float *vectors;
73 Dimension dim;
74 } CentroidEncoderHalf;
75
76 static inline void
77 88 centroid_encode_half(CentroidEncoder *enc, uint32_t index, void *dest)
78 {
79 88 CentroidEncoderHalf *he = (CentroidEncoderHalf *)enc;
80 88 const float *src = he->vectors + (size_t)index * he->dim;
81 88 vs_float_to_half_array(src, (half *)dest, he->dim);
82 88 }
83
84 static const CentroidEncoderOps centroid_encoder_half_ops = {
85 .encode_into = centroid_encode_half,
86 };
87
88 /* RaBitQ encoder */
89 typedef struct
90 {
91 CentroidEncoder base;
92 const float *vectors;
93 Dimension dim;
94 const RaBitQParams *params;
95 const float *global_mean;
96 } CentroidEncoderRaBitQ;
97
98 static inline void
99 2460 centroid_encode_rabitq(CentroidEncoder *enc, uint32_t index, void *dest)
100 {
101 2460 CentroidEncoderRaBitQ *re = (CentroidEncoderRaBitQ *)enc;
102 2460 Vec32Ref vref = {
103 2460 .data = re->vectors + (size_t)index * re->dim,
104 1820 .dim = re->dim,
105 };
106 2460 Vec32Ref mref = {.data = re->global_mean, .dim = re->dim};
107 2460 vs_rabitq_encode_into(re->params, vref, mref, (RaBitQData *)dest);
108 2460 }
109
110 static const CentroidEncoderOps centroid_encoder_rabitq_ops = {
111 .encode_into = centroid_encode_rabitq,
112 };
113
114 /* ----------------------------------------------------------------
115 * CentroidEncoderState — stack-allocated encoder storage
116 *
117 * Union of all built-in encoder structs. Callers declare one on
118 * the stack and pass it to centroid_encoder_init(), which returns
119 * a CentroidEncoder * pointing into the union.
120 *
121 * RaBitQ params/global_mean are only used when fmt is RABITQ;
122 * callers may pass NULL for unused fields.
123 * ---------------------------------------------------------------- */
124 typedef union
125 {
126 CentroidEncoderFloat f;
127 CentroidEncoderHalf h;
128 CentroidEncoderRaBitQ r;
129 } CentroidEncoderState;
130
131 static inline CentroidEncoder *
132 710 centroid_encoder_init(
133 CentroidEncoderState *state,
134 PrismCentroidFormat fmt,
135 const float *vectors,
136 Dimension dim,
137 const RaBitQParams *params,
138 const float *global_mean)
139 {
140
3/5
✓ Branch 0 taken 200 times.
✓ Branch 1 taken 10 times.
✓ Branch 2 taken 500 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
 710 switch (fmt)
141 {
142 200 case PRISM_CENTROID_FMT_FLOAT:
143 200 state->f = (CentroidEncoderFloat){
144 .base.ops = &centroid_encoder_float_ops,
145 .vectors = vectors,
146 .dim = dim,
147 };
148 200 return &state->f.base;
149 10 case PRISM_CENTROID_FMT_HALF:
150 10 state->h = (CentroidEncoderHalf){
151 .base.ops = &centroid_encoder_half_ops,
152 .vectors = vectors,
153 .dim = dim,
154 };
155 10 return &state->h.base;
156 500 case PRISM_CENTROID_FMT_RABITQ:
157 500 state->r = (CentroidEncoderRaBitQ){
158 .base.ops = &centroid_encoder_rabitq_ops,
159 .vectors = vectors,
160 .dim = dim,
161 .params = params,
162 .global_mean = global_mean,
163 };
164 500 return &state->r.base;
165 ✗ case PRISM_CENTROID_FMT_FASTSCAN:
166 /* fastscan uses a group-packed page layout that doesn't fit
167 * this per-vector encoder model. The build path emits these
168 * pages directly (see prism_centroid_write_fastscan_pages —
169 * future work) or they're produced by an in-place conversion. */
170 ✗ return NULL;
171 }
172 ✗ return NULL;
173 }
174
175 /*
176 * Write centroid entries to linked pages.
177 *
178 * Fills pages to capacity, linking via next_blkno. When start_blkno
179 * is InvalidBlockNumber, appends via new_page. Otherwise writes to
180 * pre-reserved blocks starting at start_blkno via write_page.
181 *
182 * Returns the BlockNumber of the first centroid page.
183 */
184 BlockNumber prism_centroid_write_pages(
185 VsStorage *storage,
186 Dimension dim,
187 uint32_t nlist,
188 PrismCentroidFormat fmt,
189 uint8_t level,
190 uint16_t flags,
191 uint16_t child_count,
192 CentroidEncoder *encoder,
193 const BlockNumber *child_blknos,
194 const float *pt_centroids,
195 BlockNumber start_blkno);
196
197 /*
198 * Write centroid entries as FASTSCAN-format pages.
199 *
200 * Same inputs as prism_centroid_write_pages but emits 32-vector groups
201 * with kPerm0-packed RaBitQ codes instead of per-entry RaBitQData,
202 * so the scan path uses vs_fastscan_accumulate_hacc rather than
203 * vs_rabitq_inner_product_multi at score time.
204 *
205 * Requires `params` and `global_mean` (the RaBitQ encoding inputs);
206 * the float `vectors` array provides the centroids to encode. Each
207 * entry's f_error is computed from f_add / f_rescale at emit time
208 * so the score path doesn't need to recompute it.
209 */
210 BlockNumber prism_centroid_write_fastscan_pages(
211 VsStorage *storage,
212 Dimension dim,
213 uint32_t nlist,
214 uint8_t level,
215 uint16_t flags,
216 const RaBitQParams *params,
217 const float *vectors,
218 const float *global_mean,
219 const BlockNumber *child_blknos,
220 BlockNumber start_blkno);
221
222 #endif /* PRISM_CENTROID_BUILD_H */
223