GCC Code Coverage Report

Directory: src/
File: src/algo/hkmeans.h
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 6 6 100.0%
Functions: 3 3 100.0%
Branches: 3 4 75.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 * hkmeans.h - Hierarchical k-means tree builder
6 *
7 * Builds a BFS-ordered tree of centroids using hierarchical k-means.
8 * At each node, runs vs_kmeans_f32() to split vectors into fan_out
9 * children, then recurses on each partition.
10 *
11 * The result is self-contained: all data is stored in a single
12 * contiguous allocation using byte offsets instead of pointers.
13 * This allows memcpy into shared memory (DSM) for parallel builds.
14 *
15 * Used by both the PG IAM build (build.c) and the CLI
16 * benchmark (bench_search.c).
17 */
18
19 #ifndef VS_HKMEANS_H
20 #define VS_HKMEANS_H
21
22 #include <stdint.h>
23
24 #include "algo/kmeans.h"
25 #include "core/types.h"
26
27 /*
28 * HKMeansNode - One node in the BFS-ordered tree
29 *
30 * centroid_offset is a byte offset from the HKMeansResult base,
31 * making the tree self-contained and memcpy-able.
32 */
33 typedef struct HKMeansNode
34 {
35 uint32_t centroid_offset; /* byte offset from HKMeansResult* */
36 uint32_t nchildren; /* actual cluster count (<= fan_out) */
37 uint32_t level; /* tree level (0 = root) */
38 uint32_t first_child; /* index in nodes[] of first child */
39 uint32_t first_leaf; /* offset into leaf_centroids (leaf-parent only) */
40 } HKMeansNode;
41
42 /* Sentinel for leaf nodes with no children */
43 #define HKMEANS_NO_CHILD UINT32_MAX
44
45 /*
46 * HKMeansResult - Complete hierarchical k-means tree
47 *
48 * Self-contained: all data (nodes, leaf centroids, internal
49 * centroids) is stored in a single contiguous allocation.
50 * The struct can be memcpy'd into shared memory (DSM) and
51 * used directly by other processes without deserialization.
52 *
53 * Layout:
54 * [HKMeansResult header]
55 * [HKMeansNode nodes[nnodes]] — at nodes_offset
56 * [float leaf_centroids[nleaves*dim]] — at leaf_offset
57 * [float internal_centroids[...]] — packed after leaves
58 */
59 typedef struct HKMeansResult
60 {
61 uint32_t nodes_offset; /* byte offset to nodes[] */
62 uint32_t leaf_offset; /* byte offset to leaf centroids */
63 uint32_t total_size; /* total allocation size in bytes */
64 uint32_t nnodes; /* total internal nodes */
65 uint32_t nlevels; /* tree depth */
66 uint32_t nleaves; /* total leaf centroids */
67 uint32_t fan_out; /* children per node (max) */
68 Dimension dim; /* vector dimension */
69 } HKMeansResult;
70
71 /* ----------------------------------------------------------------
72 * Accessors — resolve byte offsets to typed pointers
73 * ---------------------------------------------------------------- */
74
75 static inline HKMeansNode *
76 13601 hk_nodes(const HKMeansResult *r)
77 {
78
3/4
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 187 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 414 times.
 13601 return (HKMeansNode *)((char *)r + r->nodes_offset);
79 }
80
81 static inline float *
82 2107 hk_leaf_centroids(const HKMeansResult *r)
83 {
84 2107 return (float *)((char *)r + r->leaf_offset);
85 }
86
87 static inline float *
88 23779 hk_node_centroids(const HKMeansResult *r, const HKMeansNode *node)
89 {
90 23779 return (float *)((char *)r + node->centroid_offset);
91 }
92
93 /* ----------------------------------------------------------------
94 * Public API
95 * ---------------------------------------------------------------- */
96
97 /*
98 * Build a hierarchical k-means tree.
99 *
100 * indices: optional index array for indirect access (NULL = identity).
101 * When non-NULL, vector i is at vectors[indices[i] * dim].
102 * This allows subsampling without copying.
103 *
104 * Returns a single contiguous allocation on success, NULL on failure.
105 * Caller must free with vs_free().
106 */
107 HKMeansResult *vs_hkmeans_f32(
108 const float *vectors,
109 uint32_t nvecs,
110 const uint32_t *indices,
111 Dimension dim,
112 uint32_t nlist,
113 uint32_t fan_out,
114 DistanceMetric metric,
115 const KMeansOptions *options);
116
117 /*
118 * Route a vector to its nearest leaf centroid by descending the tree.
119 * Returns the leaf index (0..nleaves-1).
120 *
121 * Optionally writes the distance to the nearest leaf centroid into
122 * *out_distance (may be NULL).
123 */
124 uint32_t vs_hkmeans_assign(
125 const HKMeansResult *tree,
126 const float *vec,
127 DistanceMetric metric,
128 Distance *out_distance);
129
130 /* Upper bound on k / beam_width for vs_hkmeans_assign_topk (keeps the
131 * beam scratch on the stack). */
132 #define VS_HK_MAX_TOPK 64
133
134 /*
135 * Beam-search the tree for the k nearest leaf centroids.
136 *
137 * Maintains a beam of the best `beam_width` nodes per level, then keeps
138 * the k nearest leaves at the leaf level. Approximate for k/beam_width
139 * smaller than the tree fan-out, but far cheaper than scanning all
140 * leaves — used for secondary (boundary) cluster assignment during
141 * build. k and beam_width are clamped to VS_HK_MAX_TOPK.
142 *
143 * out_leaves[k] receives leaf indices sorted by ascending distance;
144 * out_dists[k] (optional) the matching distances. Returns the number of
145 * leaves written (<= k).
146 */
147 uint32_t vs_hkmeans_assign_topk(
148 const HKMeansResult *tree,
149 const float *vec,
150 DistanceMetric metric,
151 uint32_t k,
152 uint32_t beam_width,
153 uint32_t *out_leaves,
154 Distance *out_dists);
155
156 /*
157 * Upper bound (bytes) on the contiguous size of a tree built for `nlist`
158 * leaves with `fan_out`. Used to size the fixed per-subtree DSM slots the
159 * parallel build's participants write their subtrees into.
160 */
161 size_t
162 vs_hkmeans_max_blob_size(uint32_t nlist, uint32_t fan_out, Dimension dim);
163
164 /*
165 * Same bound with every level's width capped at max_leaves: a node exists
166 * only where a training vector landed, so a subtree clustered from
167 * max_leaves vectors can never exceed it, whatever the nlist target.
168 */
169 size_t vs_hkmeans_max_blob_size_capped(
170 uint32_t nlist, uint32_t fan_out, Dimension dim, uint64_t max_leaves);
171
172 /*
173 * Tree depth for `nlist` leaves at `fan_out` — the same value the tree build
174 * uses internally. Exposed so the streaming (page-backed) centroid-tree build
175 * can compute the level structure without materializing a tree.
176 */
177 uint32_t vs_hkmeans_nlevels(uint32_t nlist, uint32_t fan_out);
178
179 /*
180 * Build a one-level (flat) tree directly from pre-computed leaf centroids.
181 *
182 * For a flat clustering (nleaves <= fan_out) the root k-means already produced
183 * every leaf centroid, so the parallel build can assemble the tree straight
184 * from them rather than re-gathering the samples and re-clustering. Produces
185 * the same shape vs_hkmeans_f32 does for a single-level build: one
186 * leaf-parent root with nleaves children. Returns a contiguous allocation;
187 * caller frees with vs_free().
188 */
189 HKMeansResult *vs_hkmeans_build_flat(
190 const float *centroids,
191 uint32_t nleaves,
192 uint32_t fan_out,
193 Dimension dim);
194
195 #endif /* VS_HKMEANS_H */
196