| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | /* | ||
| 2 | * Copyright (c) 2026 Tiger Data, Inc. | ||
| 3 | * Licensed under the PostgreSQL License. See LICENSE for details. | ||
| 4 | * | ||
| 5 | * parallel_build.h - Parallel index build (worker + leader shared decls) | ||
| 6 | * | ||
| 7 | * Phased parallel build, barrier-synchronized: | ||
| 8 | * Sampling: workers cooperatively scan the heap into a bounded, | ||
| 9 | * budget-sized sample region (a dedicated segment, released once | ||
| 10 | * clustering is done). | ||
| 11 | * Clustering: root k-means over the sample (leader reduces between | ||
| 12 | * iterations), then per-root-child subtrees in batched ring slots; | ||
| 13 | * the leader records each batch's layout and spills the subtree | ||
| 14 | * blobs, then streams every centroid + head page to disk by itself. | ||
| 15 | * Refine (when subsampled): workers route the full table page-backed | ||
| 16 | * and accumulate per-leaf means; the leader rewrites the heads. | ||
| 17 | * Posting: workers route + encode into a cluster-keyed shared sort; | ||
| 18 | * the leader merges runs and writes the posting lists. | ||
| 19 | */ | ||
| 20 | |||
| 21 | #ifndef PRISM_PARALLEL_BUILD_H | ||
| 22 | #define PRISM_PARALLEL_BUILD_H | ||
| 23 | |||
| 24 | /* | ||
| 25 | * Back-end primitives the parallel build runs on. The standalone shims live in | ||
| 26 | * src/standalone/; a PostgreSQL build uses the real PG headers. This guard is | ||
| 27 | * the single place that switch is made — the standalone headers carry no PG | ||
| 28 | * branch of their own. | ||
| 29 | */ | ||
| 30 | #ifdef VS_STANDALONE | ||
| 31 | #include "standalone/barrier.h" | ||
| 32 | #include "standalone/pg_compat.h" /* Size, BlockNumber, ItemPointerData, Relation */ | ||
| 33 | #include "standalone/shm_toc.h" | ||
| 34 | #else | ||
| 35 | #include <postgres.h> | ||
| 36 | |||
| 37 | #include <storage/barrier.h> | ||
| 38 | #include <storage/block.h> | ||
| 39 | #include <storage/dsm.h> | ||
| 40 | #include <storage/itemptr.h> | ||
| 41 | #include <storage/shm_toc.h> | ||
| 42 | #include <utils/rel.h> | ||
| 43 | #endif | ||
| 44 | |||
| 45 | #include <inttypes.h> | ||
| 46 | |||
| 47 | #include "algo/hkmeans.h" | ||
| 48 | #include "core/log.h" | ||
| 49 | #include "core/memory.h" | ||
| 50 | #include "core/types.h" | ||
| 51 | #include "index/centroid_page.h" /* PrismCentroidFormat */ | ||
| 52 | #include "index/posting_build.h" | ||
| 53 | #include "index/storage.h" /* VsStorage */ | ||
| 54 | #include "quant/rabitq.h" | ||
| 55 | |||
| 56 | /* Forward decl so prism_build_scan's prototype can reference it without | ||
| 57 | * pulling in the executor headers; callers that pass one already have the full | ||
| 58 | * type. | ||
| 59 | */ | ||
| 60 | struct IndexInfo; | ||
| 61 | |||
| 62 | /* | ||
| 63 | * Per-vector build scan callback (back-end-neutral): the PG scan adapter | ||
| 64 | * unwraps each heap tuple's Datum to the vector's data pointer, and the | ||
| 65 | * standalone work-stealing scan points straight into its vector array. The tid | ||
| 66 | * is the real heap TID under PG and a reversibly-synthesized one in | ||
| 67 | * standalone. | ||
| 68 | */ | ||
| 69 | typedef void (*PrismBuildScanCb)( | ||
| 70 | void *state, ItemPointerData tid, const float *vec); | ||
| 71 | |||
| 72 | /* ---------------------------------------------------------------- | ||
| 73 | * shm_toc region keys (the toc is a DSM segment in PG, a heap arena in | ||
| 74 | * standalone). A few regions are PG-only (WAL/buffer usage, query text) but | ||
| 75 | * their keys live here with the rest for one contiguous numbering. | ||
| 76 | * ---------------------------------------------------------------- */ | ||
| 77 | |||
| 78 | #define PRISM_DSM_KEY_SHARED UINT64CONST(0xB000000000000001) | ||
| 79 | #define PRISM_DSM_KEY_WAL_USAGE UINT64CONST(0xB000000000000006) | ||
| 80 | #define PRISM_DSM_KEY_BUFFER_USAGE UINT64CONST(0xB000000000000007) | ||
| 81 | #define PRISM_DSM_KEY_QUERY_TEXT UINT64CONST(0xB000000000000008) | ||
| 82 | #define PRISM_DSM_KEY_BARRIER UINT64CONST(0xB000000000000009) | ||
| 83 | #define PRISM_DSM_KEY_SAMPLES UINT64CONST(0xB00000000000000A) | ||
| 84 | #define PRISM_DSM_KEY_CENTROIDS UINT64CONST(0xB00000000000000B) | ||
| 85 | #define PRISM_DSM_KEY_KM_WORKERS UINT64CONST(0xB00000000000000C) | ||
| 86 | #define PRISM_DSM_KEY_ROOT_ASSIGN UINT64CONST(0xB00000000000000E) | ||
| 87 | #define PRISM_DSM_KEY_SORTSHARED UINT64CONST(0xB000000000000012) | ||
| 88 | /* Page-backed phase-3 routing: the global mean, published by the leader before | ||
| 89 | * the tree-ready barrier so workers route exactly as the query/insert paths | ||
| 90 | * do. The posting-head base (leaf c's head = first_posting + c) is a scalar in | ||
| 91 | * PrismBuildShared, not a shared array. */ | ||
| 92 | #define PRISM_DSM_KEY_GLOBAL_MEAN UINT64CONST(0xB000000000000014) | ||
| 93 | |||
| 94 | /* ---------------------------------------------------------------- | ||
| 95 | * PrismBuildShared — back-end-neutral shared build state | ||
| 96 | * | ||
| 97 | * Immutable config set by the leader before the workers start, plus counters | ||
| 98 | * the workers update concurrently during the posting scan under a lock the | ||
| 99 | * back-end owns (prism_pbuild_worker_add_counts). Each back-end embeds this as | ||
| 100 | * the first member of its own struct carrying the back-end-specific state — | ||
| 101 | * for PG, the relation OIDs, query id, the spinlock, and a trailing | ||
| 102 | * ParallelTableScanDesc (see PrismBuildSharedPg in parallel_backend.c). | ||
| 103 | * ---------------------------------------------------------------- */ | ||
| 104 | |||
| 105 | typedef struct PrismBuildShared | ||
| 106 | { | ||
| 107 | /* Immutable — set by the leader before launch */ | ||
| 108 | Dimension dim; | ||
| 109 | DistanceMetric metric; | ||
| 110 | uint32_t nlist; | ||
| 111 | uint32_t fan_out; | ||
| 112 | double soar_lambda; | ||
| 113 | double boundary_epsilon; | ||
| 114 | bool fastscan; | ||
| 115 | PrismCentroidFormat centroid_format; | ||
| 116 | uint64_t rabitq_seed; | ||
| 117 | /* Planned as 1 + planned workers (it sizes the per-participant DSM | ||
| 118 | * regions), then narrowed by prism_pbuild_launch to 1 + the workers that | ||
| 119 | * actually started when the launch falls short. Work partitioned by | ||
| 120 | * participant must use this count; region sizing keeps the planned | ||
| 121 | * value and leaves the tail slots unused. */ | ||
| 122 | int nparticipants; | ||
| 123 | /* CREATE INDEX CONCURRENTLY: the leader scans with an MVCC snapshot and | ||
| 124 | * participants take weak relation locks. Workers must mark their rebuilt | ||
| 125 | * IndexInfo concurrent too, or heapam's snapshot/OldestXmin check trips. | ||
| 126 | */ | ||
| 127 | bool concurrent; | ||
| 128 | /* Posting sort work budget (KB); PG sets it from maintenance_work_mem, | ||
| 129 | * standalone leaves it 0 (its in-memory sorter ignores it). The shared | ||
| 130 | * phase-3 code reads this instead of the PG-only GUC. */ | ||
| 131 | int work_mem_kb; | ||
| 132 | |||
| 133 | /* K-means config */ | ||
| 134 | uint32_t max_samples_per_worker; | ||
| 135 | uint32_t km_max_iterations; | ||
| 136 | float km_tolerance; | ||
| 137 | uint32_t km_k; /* root k-means k (= fan_out) */ | ||
| 138 | |||
| 139 | /* Refine gate input, set at setup: the sample-per-leaf threshold below | ||
| 140 | * which a subsampled build refines the leaf encode references on the | ||
| 141 | * full table (0 = refinement off; standalone leaves it 0). Whether the | ||
| 142 | * build IS subsampled is decided empirically from the sampling pass: | ||
| 143 | * each participant counts the live rows its scan saw next to the rows | ||
| 144 | * it kept, and kept < seen means the sample excludes real rows. */ | ||
| 145 | uint32_t refine_threshold; | ||
| 146 | |||
| 147 | /* The refine decision. The LEADER makes it after clustering -- the | ||
| 148 | * actual leaf count is only known then, and dividing the collected | ||
| 149 | * samples by the worst-case leaf bound would understate samples-per-leaf | ||
| 150 | * and refine too eagerly -- and publishes it before the tree-ready | ||
| 151 | * barrier. Workers read it after that barrier, so leader and workers | ||
| 152 | * gate the refine phase (and its barriers) identically. */ | ||
| 153 | bool refine; | ||
| 154 | |||
| 155 | /* Tile capacity (leaves) of the refine accumulator that overlays the | ||
| 156 | * sample region once sampling is done; sized at setup | ||
| 157 | * (the back-end bounds it by its memory budget AND the sample region's | ||
| 158 | * size, since the overlay lives inside that region). */ | ||
| 159 | uint32_t refine_tile_cap; | ||
| 160 | |||
| 161 | /* Counters — updated concurrently under the back-end's lock */ | ||
| 162 | double reltuples; | ||
| 163 | double indtuples; | ||
| 164 | double soar_dupes; | ||
| 165 | |||
| 166 | /* K-means convergence — set by leader between barriers */ | ||
| 167 | bool km_converged; | ||
| 168 | |||
| 169 | /* Per-child subtree blob slot size (bytes) in the child-subtrees region; | ||
| 170 | * set by the leader before launch so workers can index their slot. */ | ||
| 171 | /* Written by the leader after root assignment (see the subtree-ring | ||
| 172 | * seam); workers read it after the ring barrier. */ | ||
| 173 | uint64_t subtree_slot_size; | ||
| 174 | |||
| 175 | /* Page-backed routing knobs (mirror the prism.centroid_* GUCs), so | ||
| 176 | * phase-3 workers build a PrismIndexBase that routes identically to | ||
| 177 | * the query path. fastscan_bits is the FASTSCAN centroid bit width | ||
| 178 | * (base.fastscan when the centroid format is FASTSCAN; 0 otherwise). */ | ||
| 179 | float centroid_error_scale; | ||
| 180 | float centroid_beam_scale; | ||
| 181 | int fastscan_bits; | ||
| 182 | |||
| 183 | /* Published by the leader after the streaming tree write: the | ||
| 184 | * centroid-tree root block (workers' phase-3 | ||
| 185 | * PrismIndexBase.first_centroid) and the tree depth (base.nlevels). The | ||
| 186 | * tree itself lives only on pages. | ||
| 187 | */ | ||
| 188 | BlockNumber first_centroid; | ||
| 189 | uint8_t nlevels; | ||
| 190 | |||
| 191 | /* Published by the leader before phase 3: the first posting-head block. | ||
| 192 | * Cluster c's head is first_posting + c (formula), so workers map a routed | ||
| 193 | * head block back to its leaf by subtraction — no O(nlist) head array. */ | ||
| 194 | BlockNumber first_posting; | ||
| 195 | } PrismBuildShared; | ||
| 196 | |||
| 197 | /* ---------------------------------------------------------------- | ||
| 198 | * K-means sampling: per-worker sample slots in DSM | ||
| 199 | * | ||
| 200 | * Layout: [sample_counts[nparticipants]] [rows_seen[nparticipants]] then | ||
| 201 | * [samples[worker_id][max_per_worker * dim]] packed. | ||
| 202 | * ---------------------------------------------------------------- */ | ||
| 203 | |||
| 204 | typedef struct PrismDsmSamples | ||
| 205 | { | ||
| 206 | uint32_t nparticipants; | ||
| 207 | uint32_t max_per_worker; | ||
| 208 | Dimension dim; | ||
| 209 | /* counts[nparticipants] (rows kept), then MAXALIGN'd | ||
| 210 | * seen[nparticipants] (live rows the sampling scan visited; kept < seen | ||
| 211 | * means the sample is a strict subset of the table -- 64-bit, since a | ||
| 212 | * participant's scan share is unbounded and a wrapped counter would | ||
| 213 | * silently flip the refine gate), then the per-worker sample blocks | ||
| 214 | * (samples[worker][max_per_worker * dim]) packed right after. */ | ||
| 215 | uint32_t counts[]; | ||
| 216 | } PrismDsmSamples; | ||
| 217 | |||
| 218 | static inline uint32_t * | ||
| 219 | 4702 | prism_dsm_sample_counts(PrismDsmSamples *s) | |
| 220 | { | ||
| 221 |
4/4✓ Branch 0 taken 576 times.
✓ Branch 1 taken 561 times.
✓ Branch 2 taken 44 times.
✓ Branch 3 taken 86 times.
|
4572 | return s->counts; |
| 222 | } | ||
| 223 | |||
| 224 | static inline uint64_t * | ||
| 225 | 2618 | prism_dsm_sample_seen(PrismDsmSamples *s) | |
| 226 | { | ||
| 227 | 1557 | return (uint64_t *)MAXALIGN(s->counts + s->nparticipants); | |
| 228 | } | ||
| 229 | |||
| 230 | static inline float * | ||
| 231 | 1969 | prism_dsm_worker_samples(PrismDsmSamples *s, int worker_id) | |
| 232 | { | ||
| 233 | /* The sample blocks begin after counts[] and seen[]. */ | ||
| 234 | 1969 | float *samples = (float *)(prism_dsm_sample_seen(s) + s->nparticipants); | |
| 235 | 1969 | return samples + (size_t)worker_id * s->max_per_worker * s->dim; | |
| 236 | } | ||
| 237 | |||
| 238 | static inline Size | ||
| 239 | 348 | prism_dsm_samples_size( | |
| 240 | int nparticipants, uint32_t max_per_worker, Dimension dim) | ||
| 241 | { | ||
| 242 | 348 | Size sz = offsetof(PrismDsmSamples, counts); | |
| 243 | 348 | sz += (Size)nparticipants * sizeof(uint32_t); | |
| 244 | 348 | sz = MAXALIGN(sz); | |
| 245 | 348 | sz += (Size)nparticipants * sizeof(uint64_t); | |
| 246 | 348 | sz += (Size)nparticipants * max_per_worker * dim * sizeof(float); | |
| 247 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 45 times.
|
348 | return sz; |
| 248 | } | ||
| 249 | |||
| 250 | /* ---------------------------------------------------------------- | ||
| 251 | * Root assignments: per-worker uint32_t[max_per_worker] in DSM | ||
| 252 | * | ||
| 253 | * After root k-means converges, each worker stores its root | ||
| 254 | * assignments here. Used to filter samples during child k-means. | ||
| 255 | * ---------------------------------------------------------------- */ | ||
| 256 | |||
| 257 | typedef struct PrismDsmRootAssign | ||
| 258 | { | ||
| 259 | uint32_t nparticipants; | ||
| 260 | uint32_t max_per_worker; | ||
| 261 | /* assignments[worker][max_per_worker] packed contiguously. */ | ||
| 262 | uint32_t assignments[]; | ||
| 263 | } PrismDsmRootAssign; | ||
| 264 | |||
| 265 | static inline uint32_t * | ||
| 266 | 1618 | prism_dsm_root_assignments(PrismDsmRootAssign *ra, int worker_id) | |
| 267 | { | ||
| 268 | 1618 | return ra->assignments + (size_t)worker_id * ra->max_per_worker; | |
| 269 | } | ||
| 270 | |||
| 271 | static inline Size | ||
| 272 | 262 | prism_dsm_root_assign_size(int nparticipants, uint32_t max_per_worker) | |
| 273 | { | ||
| 274 | 262 | Size sz = offsetof(PrismDsmRootAssign, assignments); | |
| 275 | 262 | sz += (Size)nparticipants * max_per_worker * sizeof(uint32_t); | |
| 276 | 262 | return sz; | |
| 277 | } | ||
| 278 | |||
| 279 | /* ---------------------------------------------------------------- | ||
| 280 | * K-means shared centroids in DSM | ||
| 281 | * | ||
| 282 | * centroids[nlist * dim] + norms_c[nlist] | ||
| 283 | * Written by leader between barriers, read by workers during | ||
| 284 | * assignment. | ||
| 285 | * ---------------------------------------------------------------- */ | ||
| 286 | |||
| 287 | static inline Size | ||
| 288 | 217 | prism_dsm_centroids_size(uint32_t nlist, Dimension dim) | |
| 289 | { | ||
| 290 | 217 | return (Size)nlist * dim * sizeof(float) + (Size)nlist * sizeof(float); | |
| 291 | } | ||
| 292 | |||
| 293 | static inline float * | ||
| 294 | 980 | prism_dsm_centroids(char *base) | |
| 295 | { | ||
| 296 | 980 | return (float *)base; | |
| 297 | } | ||
| 298 | |||
| 299 | /* ---------------------------------------------------------------- | ||
| 300 | * Child subtrees: per-root-child HKMeansResult blobs in DSM | ||
| 301 | * | ||
| 302 | * After the root k-means splits the samples into fan_out groups, each group's | ||
| 303 | * subtree is built independently and written, as a contiguous HKMeansResult, | ||
| 304 | * into a fixed-size slot. The batched streaming build keeps only a bounded | ||
| 305 | * ring of nparticipants slots resident (each participant owns slot | ||
| 306 | * participant_id; the leader streams each batch to pages before the next batch | ||
| 307 | * reuses the ring), so the region is O(nparticipants * slot_size), independent | ||
| 308 | * of the partition count. The leader sizes the slot only after root | ||
| 309 | * assignment -- a subtree can never hold more leaves than its child's | ||
| 310 | * samples, which caps the slot far below the analytic worst case -- and | ||
| 311 | * creates the ring as its own segment then (the subtree-ring seam below); | ||
| 312 | * the slot size travels in PrismBuildShared.subtree_slot_size. | ||
| 313 | * ---------------------------------------------------------------- */ | ||
| 314 | |||
| 315 | static inline Size | ||
| 316 | 52 | prism_dsm_child_subtrees_size(uint32_t nslots, uint64_t slot_size) | |
| 317 | { | ||
| 318 | 52 | return (Size)nslots * (Size)slot_size; | |
| 319 | } | ||
| 320 | |||
| 321 | static inline char * | ||
| 322 | 388 | prism_dsm_child_subtree(char *base, uint32_t slot, uint64_t slot_size) | |
| 323 | { | ||
| 324 | 388 | return base + (Size)slot * (Size)slot_size; | |
| 325 | } | ||
| 326 | |||
| 327 | /* | ||
| 328 | * Worst-case leaf count for a tree targeting `nlist` leaves at this fan_out: | ||
| 329 | * fan_out^nlevels (clamped to >= nlist). The exact leaf count isn't known | ||
| 330 | * until k-means runs, so the parallel build uses this to size its DSM regions | ||
| 331 | * up front, then narrows to the real tree->nleaves afterward. | ||
| 332 | * | ||
| 333 | * Accumulating in 64 bits keeps a product past UINT32_MAX an error rather | ||
| 334 | * than a truncation: every DSM region is sized from this, so a wrapped count | ||
| 335 | * describes a smaller tree than the one being built. The reloption bounds | ||
| 336 | * keep the product under 2^32 today, which is what makes the ceiling a | ||
| 337 | * guard on those bounds rather than a limit callers meet. | ||
| 338 | */ | ||
| 339 | static inline uint32_t | ||
| 340 | 141 | prism_max_nlist(uint32_t nlist, uint32_t fan_out) | |
| 341 | { | ||
| 342 | 141 | uint32_t nlevels = vs_hkmeans_nlevels(nlist, fan_out); | |
| 343 | 141 | uint64_t max_nlist = 1; | |
| 344 |
3/3✓ Branch 0 taken 192 times.
✓ Branch 1 taken 168 times.
✓ Branch 2 taken 45 times.
|
405 | for (uint32_t l = 0; l < nlevels; l++) |
| 345 | { | ||
| 346 | 264 | max_nlist *= fan_out; | |
| 347 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 264 times.
|
264 | if (max_nlist > UINT32_MAX) |
| 348 |
0/2✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
72 | vs_error( |
| 349 | "prism: %u levels at fan_out %u reach %" PRIu64 | ||
| 350 | " leaves, past the largest addressable count " | ||
| 351 | "(nlist %u); decrease fan_out or nlist", | ||
| 352 | nlevels, | ||
| 353 | fan_out, | ||
| 354 | max_nlist, | ||
| 355 | nlist); | ||
| 356 | } | ||
| 357 |
1/2✓ Branch 0 taken 141 times.
✗ Branch 1 not taken.
|
141 | return max_nlist < nlist ? nlist : (uint32_t)max_nlist; |
| 358 | } | ||
| 359 | |||
| 360 | static inline float * | ||
| 361 | 652 | prism_dsm_norms_c(char *base, uint32_t nlist, Dimension dim) | |
| 362 | { | ||
| 363 | /* norms_c[nlist] sits right after centroids[nlist * dim]. */ | ||
| 364 |
2/2✓ Branch 1 taken 44 times.
✓ Branch 2 taken 86 times.
|
652 | return prism_dsm_centroids(base) + (size_t)nlist * dim; |
| 365 | } | ||
| 366 | |||
| 367 | /* ---------------------------------------------------------------- | ||
| 368 | * K-means per-worker accumulators in DSM | ||
| 369 | * | ||
| 370 | * Per worker: centroid_sums[nlist * dim] + centroid_cnts[nlist] | ||
| 371 | * + cost (1 float) | ||
| 372 | * ---------------------------------------------------------------- */ | ||
| 373 | |||
| 374 | static inline Size | ||
| 375 | 3407 | prism_dsm_km_worker_size(uint32_t nlist, Dimension dim) | |
| 376 | { | ||
| 377 | 3407 | return (Size)nlist * dim * sizeof(float) + (Size)nlist * sizeof(uint32_t) + | |
| 378 | sizeof(float); | ||
| 379 | } | ||
| 380 | |||
| 381 | static inline Size | ||
| 382 | 543 | prism_dsm_km_workers_size(int nparticipants, uint32_t nlist, Dimension dim) | |
| 383 | { | ||
| 384 | 588 | return (Size)nparticipants * prism_dsm_km_worker_size(nlist, dim); | |
| 385 | } | ||
| 386 | |||
| 387 | static inline float * | ||
| 388 | 4629 | prism_dsm_km_worker_sums( | |
| 389 | char *base, uint32_t nlist, Dimension dim, int worker_id) | ||
| 390 | { | ||
| 391 | 7493 | return (float *)(base + | |
| 392 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 86 times.
|
3124 | (size_t)worker_id * prism_dsm_km_worker_size(nlist, dim)); |
| 393 | } | ||
| 394 | |||
| 395 | static inline uint32_t * | ||
| 396 | 4237 | prism_dsm_km_worker_cnts( | |
| 397 | char *base, uint32_t nlist, Dimension dim, int worker_id) | ||
| 398 | { | ||
| 399 | /* cnts[nlist] sits right after sums[nlist * dim]. */ | ||
| 400 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 86 times.
|
2732 | return (uint32_t *)(prism_dsm_km_worker_sums(base, nlist, dim, worker_id) + |
| 401 | 2472 | (size_t)nlist * dim); | |
| 402 | } | ||
| 403 | |||
| 404 | static inline float * | ||
| 405 | 3715 | prism_dsm_km_worker_cost( | |
| 406 | char *base, uint32_t nlist, Dimension dim, int worker_id) | ||
| 407 | { | ||
| 408 | /* cost sits right after cnts[nlist]. */ | ||
| 409 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 86 times.
|
3715 | return (float *)(prism_dsm_km_worker_cnts(base, nlist, dim, worker_id) + |
| 410 | nlist); | ||
| 411 | } | ||
| 412 | |||
| 413 | /* ---------------------------------------------------------------- | ||
| 414 | * Per-worker posting output in DSM | ||
| 415 | * | ||
| 416 | * active[nlist] per worker, packed contiguously: a flag per cluster marking | ||
| 417 | * which lists this worker wrote into. | ||
| 418 | * ---------------------------------------------------------------- */ | ||
| 419 | |||
| 420 | static inline Size | ||
| 421 | prism_dsm_worker_output_size(uint32_t nlist, int nparticipants) | ||
| 422 | { | ||
| 423 | return (Size)nparticipants * nlist * sizeof(bool); | ||
| 424 | } | ||
| 425 | |||
| 426 | static inline bool * | ||
| 427 | prism_dsm_worker_active(char *base, uint32_t nlist, int worker_id) | ||
| 428 | { | ||
| 429 | return (bool *)(base + (Size)worker_id * nlist * sizeof(bool)); | ||
| 430 | } | ||
| 431 | |||
| 432 | /* ---------------------------------------------------------------- | ||
| 433 | * Shared callbacks — used by both leader and workers | ||
| 434 | * ---------------------------------------------------------------- */ | ||
| 435 | |||
| 436 | typedef struct SampleCbState | ||
| 437 | { | ||
| 438 | float *samples; | ||
| 439 | uint32_t count; | ||
| 440 | uint64_t seen; | ||
| 441 | uint32_t max_samples; | ||
| 442 | uint32_t stride; | ||
| 443 | uint32_t stride_counter; | ||
| 444 | Dimension dim; | ||
| 445 | DistanceMetric metric; | ||
| 446 | } SampleCbState; | ||
| 447 | |||
| 448 | /* | ||
| 449 | * Shared sampling logic, called per live tuple with a raw vector pointer (no | ||
| 450 | * Datum) so the same code serves both back-ends. prism_build_scan feeds it: | ||
| 451 | * the PG scan unwraps each heap tuple's Datum, the standalone scan passes its | ||
| 452 | * in-memory vectors directly. | ||
| 453 | */ | ||
| 454 | extern void | ||
| 455 | prism_sample_cb(void *state, ItemPointerData tid, const float *vec); | ||
| 456 | |||
| 457 | extern void prism_km_assign_and_accumulate( | ||
| 458 | const float *samples, | ||
| 459 | uint32_t nsamples, | ||
| 460 | const float *centroids, | ||
| 461 | const float *norms_c, | ||
| 462 | uint32_t nlist, | ||
| 463 | Dimension dim, | ||
| 464 | DistanceMetric metric, | ||
| 465 | float *out_sums, | ||
| 466 | uint32_t *out_cnts, | ||
| 467 | float *out_cost); | ||
| 468 | |||
| 469 | /* | ||
| 470 | * Build ONE root-child's subtree into `slot`. Used by the batched streaming | ||
| 471 | * build, where each participant builds one child per batch into a ring slot | ||
| 472 | * indexed by participant (so only nparticipants subtrees are resident) and the | ||
| 473 | * leader streams each to pages. | ||
| 474 | */ | ||
| 475 | extern void prism_build_child_subtree( | ||
| 476 | uint32_t child, | ||
| 477 | uint32_t child_count, | ||
| 478 | int nparticipants, | ||
| 479 | PrismDsmSamples *dsm_samples, | ||
| 480 | PrismDsmRootAssign *dsm_ra, | ||
| 481 | const float *root_cents, | ||
| 482 | uint32_t nlist, | ||
| 483 | uint32_t fan_out, | ||
| 484 | Dimension dim, | ||
| 485 | DistanceMetric metric, | ||
| 486 | uint32_t km_max_iterations, | ||
| 487 | char *slot, | ||
| 488 | uint64_t slot_size); | ||
| 489 | |||
| 490 | /* | ||
| 491 | * Per-batch leader callback for the batched subtree stream. Fired only on the | ||
| 492 | * leader (participant 0), between the two per-batch barriers, so it can read | ||
| 493 | * the batch's finished subtrees from the slots [0, batch_size) before they are | ||
| 494 | * reused. children[s] is the child id whose subtree sits in slot s (the batch | ||
| 495 | * schedule is largest-first, not sequential). | ||
| 496 | */ | ||
| 497 | typedef void (*PrismBatchCb)( | ||
| 498 | void *arg, | ||
| 499 | const uint32_t *children, | ||
| 500 | uint32_t batch_size, | ||
| 501 | char *subtrees_base, | ||
| 502 | uint64_t slot_size); | ||
| 503 | |||
| 504 | /* | ||
| 505 | * Batched subtree build shared by the leader (participant 0) and workers: | ||
| 506 | * in ceil(km_k / nparticipants) batches, each participant clusters one root | ||
| 507 | * child's subtree into its ring slot, then (leader only) batch_cb consumes | ||
| 508 | * the batch -- recording its layout and spilling the blobs for the | ||
| 509 | * leader-only streaming write that follows; two barriers per batch keep all | ||
| 510 | * participants in lockstep. Children are scheduled largest-first (LPT, by | ||
| 511 | * root-assigned sample count with an id tie-break): every batch runs at the | ||
| 512 | * pace of its slowest subtree, so the skewed children go into the full | ||
| 513 | * batches and the tail batch gets the small ones. The schedule derives from | ||
| 514 | * shared state, so every participant computes it identically with no | ||
| 515 | * coordination; out_child_order (leader passes a km_k buffer, workers NULL) | ||
| 516 | * returns it so the blob replay can place each subtree at its child's | ||
| 517 | * reserved range. Called identically by leader and workers (workers pass | ||
| 518 | * batch_cb = NULL), so the barrier sequence matches by construction; one | ||
| 519 | * invocation per build. | ||
| 520 | */ | ||
| 521 | extern void prism_pbuild_stream_subtrees( | ||
| 522 | int participant_id, | ||
| 523 | int nparticipants, | ||
| 524 | PrismDsmSamples *dsm_samples, | ||
| 525 | PrismDsmRootAssign *dsm_ra, | ||
| 526 | const float *root_cents, | ||
| 527 | uint32_t km_k, | ||
| 528 | uint32_t nlist, | ||
| 529 | uint32_t fan_out, | ||
| 530 | Dimension dim, | ||
| 531 | DistanceMetric metric, | ||
| 532 | uint32_t km_max_iterations, | ||
| 533 | char *subtrees_base, | ||
| 534 | uint64_t slot_size, | ||
| 535 | Barrier *barrier, | ||
| 536 | PrismBatchCb batch_cb, | ||
| 537 | void *cb_arg, | ||
| 538 | uint32_t *out_child_order); | ||
| 539 | |||
| 540 | /* | ||
| 541 | * Per-participant execution of phases 1, 2, and 2b, shared by the leader | ||
| 542 | * (participant 0) and the workers — see the contract on the definitions in | ||
| 543 | * parallel_build_worker.c. Each runs its phase's per-participant work and the | ||
| 544 | * phase barrier(s); the leader-only seed/reduce in the k-means phase is gated | ||
| 545 | * on participant_id == 0. prism_pbuild_exec_kmeans returns the iteration count | ||
| 546 | * (for the leader's timing log). | ||
| 547 | */ | ||
| 548 | extern void prism_pbuild_exec_sampling( | ||
| 549 | int participant_id, | ||
| 550 | Relation heap, | ||
| 551 | Relation index, | ||
| 552 | struct IndexInfo *index_info, | ||
| 553 | PrismBuildShared *shared, | ||
| 554 | PrismDsmSamples *dsm_samples, | ||
| 555 | Barrier *barrier); | ||
| 556 | |||
| 557 | extern uint32_t prism_pbuild_exec_kmeans( | ||
| 558 | int participant_id, | ||
| 559 | PrismBuildShared *shared, | ||
| 560 | PrismDsmSamples *dsm_samples, | ||
| 561 | char *centroids_base, | ||
| 562 | char *km_workers_base, | ||
| 563 | Barrier *barrier); | ||
| 564 | |||
| 565 | extern void prism_pbuild_exec_root_assign( | ||
| 566 | int participant_id, | ||
| 567 | PrismBuildShared *shared, | ||
| 568 | PrismDsmSamples *dsm_samples, | ||
| 569 | PrismDsmRootAssign *dsm_ra, | ||
| 570 | char *centroids_base, | ||
| 571 | Barrier *barrier); | ||
| 572 | |||
| 573 | /* ---------------------------------------------------------------- | ||
| 574 | * Full-table leaf-centroid refinement (parallel) | ||
| 575 | * | ||
| 576 | * A single shared accumulator (sums[nleaves*dim] + counts[nleaves]) holds the | ||
| 577 | * per-leaf running mean for all participants -- one copy, independent of the | ||
| 578 | * worker count, so it scales to fine nlist where per-worker accumulators would | ||
| 579 | * not. Concurrent updates are guarded by a striped lock array owned by the | ||
| 580 | * back-end (prism_pbuild_accum_lock/unlock); contention is low because rows | ||
| 581 | * spread across nleaves leaves. sums accumulate in double so the refined | ||
| 582 | * means match the serial build bit-for-bit at any table size. | ||
| 583 | * ---------------------------------------------------------------- */ | ||
| 584 | |||
| 585 | #define PRISM_REFINE_LOCK_STRIPES 256 | ||
| 586 | |||
| 587 | typedef struct PrismDsmRefineAccum | ||
| 588 | { | ||
| 589 | uint32_t nleaves; | ||
| 590 | Dimension dim; | ||
| 591 | /* double sums[nleaves * dim], then uint64 counts[nleaves], packed after. | ||
| 592 | */ | ||
| 593 | double sums[FLEXIBLE_ARRAY_MEMBER]; | ||
| 594 | } PrismDsmRefineAccum; | ||
| 595 | |||
| 596 | static inline double * | ||
| 597 | 6 | prism_dsm_refine_sums(PrismDsmRefineAccum *a) | |
| 598 | { | ||
| 599 | 6 | return a->sums; | |
| 600 | } | ||
| 601 | |||
| 602 | static inline uint64_t * | ||
| 603 | 6 | prism_dsm_refine_counts(PrismDsmRefineAccum *a) | |
| 604 | { | ||
| 605 | 6 | return (uint64_t *)(a->sums + (size_t)a->nleaves * a->dim); | |
| 606 | } | ||
| 607 | |||
| 608 | static inline Size | ||
| 609 | prism_dsm_refine_accum_size(uint32_t nleaves, Dimension dim) | ||
| 610 | { | ||
| 611 | Size sz = offsetof(PrismDsmRefineAccum, sums); | ||
| 612 | sz += (Size)nleaves * dim * sizeof(double); | ||
| 613 | sz += (Size)nleaves * sizeof(uint64_t); | ||
| 614 | return sz; | ||
| 615 | } | ||
| 616 | |||
| 617 | /* | ||
| 618 | * Leaves per refine tile: the accumulator holds at most this many leaves, so | ||
| 619 | * it is a bounded constant (cap_bytes, derived from maintenance_work_mem and | ||
| 620 | * MaxAllocSize by the caller) rather than O(nleaves). nleaves above it just | ||
| 621 | * means more re-scanned tiles, not a bigger allocation. The DSM region is | ||
| 622 | * sized for this (capacity = accum->nleaves); the leader and workers derive | ||
| 623 | * the tile count from it identically, so they stay in barrier lockstep. | ||
| 624 | */ | ||
| 625 | static inline uint32_t | ||
| 626 | 48 | prism_refine_tile_leaves(uint32_t nleaves, Dimension dim, uint64_t cap_bytes) | |
| 627 | { | ||
| 628 | 48 | uint64_t per_leaf = (uint64_t)dim * sizeof(double) + sizeof(uint64_t); | |
| 629 | 48 | uint64_t t = cap_bytes / (per_leaf ? per_leaf : 1); | |
| 630 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 48 times.
|
48 | if (t < 1) |
| 631 | ✗ | t = 1; | |
| 632 | 48 | if (t > nleaves) | |
| 633 | t = nleaves; | ||
| 634 | 48 | return (uint32_t)t; | |
| 635 | } | ||
| 636 | |||
| 637 | /* | ||
| 638 | * The per-refined-leaf head writer is the shared PrismLeafWriteFn from | ||
| 639 | * posting_build.h: the leader rewrites each leaf's posting-list head with | ||
| 640 | * the full-table pt_centroid; workers pass NULL (they never divide/write). | ||
| 641 | */ | ||
| 642 | |||
| 643 | /* | ||
| 644 | * Refine the leaf encode references on the whole table, page-backed (one | ||
| 645 | * pass; routing reads only the centroid pages, which refine never rewrites, | ||
| 646 | * so the assignment is a fixed point): every row is routed exactly as the | ||
| 647 | * query/insert do (prism_query_route k=1 over the centroid pages, head -> | ||
| 648 | * leaf), per-leaf means | ||
| 649 | * accumulate into the tiled DSM accumulator, and the leader rewrites each | ||
| 650 | * leaf's head-page pt_centroid to the full-table mean. Both leader | ||
| 651 | * (participant 0) and workers call it; gated by shared->refine so they | ||
| 652 | * run the same barriers. The workers route with their own page-backed qs; the | ||
| 653 | * leader passes qs == NULL (it does not scan) and a write_head callback. | ||
| 654 | * Leaves are processed in tiles of accum->nleaves, so the accumulator stays | ||
| 655 | * bounded regardless of nlist. | ||
| 656 | */ | ||
| 657 | extern void prism_pbuild_exec_refine_paged( | ||
| 658 | int participant_id, | ||
| 659 | Relation heap, | ||
| 660 | Relation index, | ||
| 661 | struct IndexInfo *index_info, | ||
| 662 | PrismBuildShared *shared, | ||
| 663 | struct PrismQueryState *qs, | ||
| 664 | BlockNumber first_posting, | ||
| 665 | PrismDsmRefineAccum *accum, | ||
| 666 | Barrier *barrier, | ||
| 667 | PrismLeafWriteFn write_head, | ||
| 668 | void *write_head_ctx); | ||
| 669 | |||
| 670 | /* Striped lock seam for the refine accumulator (back-end owns the locks). */ | ||
| 671 | extern void prism_pbuild_accum_lock(PrismBuildShared *shared, uint32_t stripe); | ||
| 672 | extern void | ||
| 673 | prism_pbuild_accum_unlock(PrismBuildShared *shared, uint32_t stripe); | ||
| 674 | |||
| 675 | /* | ||
| 676 | * Scan every vector cooperatively, invoking cb per live tuple. Back-end seam: | ||
| 677 | * the PG implementation (parallel_backend_pg.c) drives a parallel heap scan | ||
| 678 | * and unwraps each tuple; the standalone implementation iterates its in-memory | ||
| 679 | * vector array. allow_sync/progress are PG table_index_build_scan flags, | ||
| 680 | * ignored in standalone. Returns the number of heap tuples this participant | ||
| 681 | * scanned; the posting pass accumulates those into shared->reltuples so the | ||
| 682 | * build reports the table's true row count (index_update_stats writes it to | ||
| 683 | * pg_class.reltuples). | ||
| 684 | */ | ||
| 685 | extern double prism_build_scan( | ||
| 686 | Relation heap, | ||
| 687 | Relation index, | ||
| 688 | struct IndexInfo *indexInfo, | ||
| 689 | PrismBuildShared *shared, | ||
| 690 | bool allow_sync, | ||
| 691 | bool progress, | ||
| 692 | PrismBuildScanCb cb, | ||
| 693 | void *state); | ||
| 694 | |||
| 695 | /* ---------------------------------------------------------------- | ||
| 696 | * Worker entry point — registered with CreateParallelContext | ||
| 697 | * ---------------------------------------------------------------- */ | ||
| 698 | |||
| 699 | extern void prism_parallel_build_main(dsm_segment *seg, shm_toc *toc); | ||
| 700 | |||
| 701 | /* ---------------------------------------------------------------- | ||
| 702 | * Worker lifecycle seam — back-end-specific (parallel_backend.c for PG) | ||
| 703 | * | ||
| 704 | * Runtime handles for one participant. prism_pbuild_worker_attach fills it | ||
| 705 | * when the worker joins the build; prism_pbuild_worker_detach tears it down. | ||
| 706 | * The PG version opens the heap/index relations and reports instrumentation; | ||
| 707 | * the standalone version takes the shared state and vectors directly and joins | ||
| 708 | * a thread barrier. | ||
| 709 | * ---------------------------------------------------------------- */ | ||
| 710 | |||
| 711 | typedef struct PrismPBuildWorker | ||
| 712 | { | ||
| 713 | PrismBuildShared *shared; | ||
| 714 | Barrier *barrier; | ||
| 715 | Relation heapRel; | ||
| 716 | Relation indexRel; | ||
| 717 | int worker_id; | ||
| 718 | Dimension dim; | ||
| 719 | } PrismPBuildWorker; | ||
| 720 | |||
| 721 | extern void prism_pbuild_worker_attach(shm_toc *toc, PrismPBuildWorker *w); | ||
| 722 | extern void prism_pbuild_worker_detach(shm_toc *toc, PrismPBuildWorker *w); | ||
| 723 | |||
| 724 | /* | ||
| 725 | * Page-backed phase-3 routing storage seam — back-end-specific. Workers read | ||
| 726 | * centroid + posting head pages while routing, so each needs a VsStorage over | ||
| 727 | * the index. PG (separate process) opens one on the worker's indexRel; | ||
| 728 | * standalone (threads) returns the leader's shared in-memory store published | ||
| 729 | * via prism_pbuild_publish_storage before launch. Release is a no-op for | ||
| 730 | * standalone. | ||
| 731 | */ | ||
| 732 | extern void | ||
| 733 | prism_pbuild_publish_storage(PrismBuildShared *shared, VsStorage *s); | ||
| 734 | extern VsStorage *prism_pbuild_worker_storage(PrismPBuildWorker *w); | ||
| 735 | extern void prism_pbuild_worker_storage_release(VsStorage *s); | ||
| 736 | |||
| 737 | /* | ||
| 738 | * Accumulate one worker's tuple counts into the shared state under the | ||
| 739 | * back-end's lock (the lock lives in the back-end's derived shared struct). | ||
| 740 | * heap_tuples is the participant's posting-scan share of the heap; the sum | ||
| 741 | * across participants is the table's row count, returned to PostgreSQL as | ||
| 742 | * IndexBuildResult.heap_tuples. | ||
| 743 | */ | ||
| 744 | extern void prism_pbuild_worker_add_counts( | ||
| 745 | PrismBuildShared *shared, | ||
| 746 | double indtuples, | ||
| 747 | double soar_dupes, | ||
| 748 | double heap_tuples); | ||
| 749 | |||
| 750 | /* ---------------------------------------------------------------- | ||
| 751 | * Leader launch/teardown seam — back-end-specific (parallel_backend.c for PG) | ||
| 752 | * | ||
| 753 | * prism_pbuild_launch starts the workers and blocks until the whole party has | ||
| 754 | * attached to the phase barrier (returning false, after teardown, if none | ||
| 755 | * started). Fewer workers can start than were planned (the launch competes | ||
| 756 | * for the max_parallel_workers pool), so launch narrows | ||
| 757 | * shared->nparticipants to the party that actually attached. Every phase | ||
| 758 | * that partitions work by participant reads the count after at least one | ||
| 759 | * barrier, which orders the narrowing write ahead of the read. | ||
| 760 | * prism_pbuild_teardown frees the parallel context. The standalone versions | ||
| 761 | * spawn/join threads and free the shared arena. (struct ParallelContext is | ||
| 762 | * PostgreSQL's; standalone provides its own definition.) | ||
| 763 | * ---------------------------------------------------------------- */ | ||
| 764 | |||
| 765 | struct ParallelContext; | ||
| 766 | struct WalUsage; | ||
| 767 | struct BufferUsage; | ||
| 768 | |||
| 769 | extern void prism_pbuild_teardown(struct ParallelContext *pcxt); | ||
| 770 | |||
| 771 | /* | ||
| 772 | * Sample-region seam. The k-means sample is the build's largest working set | ||
| 773 | * (up to the whole memory budget), but it is dead once the subtrees are | ||
| 774 | * clustered -- long before the posting sort claims its own budget. The PG | ||
| 775 | * back-end therefore keeps it in a dedicated DSM segment handed back through | ||
| 776 | * this seam right after the refine pass (the refine accumulator overlays the | ||
| 777 | * then-dead sample region, so it rides along for free), keeping the build's | ||
| 778 | * peak at one budget instead of stacking sample + accumulator + sort. | ||
| 779 | * Workers attach at startup and every participant releases independently; | ||
| 780 | * the segment is destroyed with the last detach. The standalone back-end | ||
| 781 | * keeps the samples in its arena (it does not bound memory) and treats | ||
| 782 | * release as a no-op. | ||
| 783 | */ | ||
| 784 | extern PrismDsmSamples *prism_pbuild_samples_attach( | ||
| 785 | shm_toc *toc, PrismBuildShared *shared, void **seg_out); | ||
| 786 | extern void prism_pbuild_samples_release(PrismDsmSamples *samples, void *seg); | ||
| 787 | |||
| 788 | /* Subtree-ring seam: leader creates the ring (own segment) after root | ||
| 789 | * assignment and publishes the slot size; workers attach after the ring | ||
| 790 | * barrier; everyone releases when the subtree batches are done. */ | ||
| 791 | extern char *prism_pbuild_subtree_ring_create( | ||
| 792 | PrismBuildShared *shared, | ||
| 793 | int nparticipants, | ||
| 794 | uint64_t slot_size, | ||
| 795 | void **seg_out); | ||
| 796 | extern char * | ||
| 797 | prism_pbuild_subtree_ring_attach(PrismBuildShared *shared, void **seg_out); | ||
| 798 | extern void prism_pbuild_subtree_ring_release(void *seg); | ||
| 799 | |||
| 800 | /* Exact-internal-centroids seam: the leader serializes the exact | ||
| 801 | * centroid collection (prism_exact_centroid_collection_write, a few MB | ||
| 802 | * even at very large nlist) it gathered while streaming the tree, and | ||
| 803 | * publishes it before the tree-ready barrier; workers attach after that | ||
| 804 | * barrier and hook it into their phase-2.5/3 routing base | ||
| 805 | * (PrismIndexBase.exact_internal). PG uses a dedicated DSM segment whose | ||
| 806 | * handle travels in the back-end shared struct; standalone shares the | ||
| 807 | * leader's allocation. Every participant releases when its routing is | ||
| 808 | * done; the segment dies with the last detach. */ | ||
| 809 | extern char *prism_pbuild_exact_centroids_create( | ||
| 810 | PrismBuildShared *shared, uint64_t nbytes, void **seg_out); | ||
| 811 | extern char * | ||
| 812 | prism_pbuild_exact_centroids_attach(PrismBuildShared *shared, void **seg_out); | ||
| 813 | extern void prism_pbuild_exact_centroids_release(void *seg); | ||
| 814 | |||
| 815 | /* One histogram over the shared root assignments: out_counts[km_k] = how | ||
| 816 | * many samples landed in each root child. Every participant derives the | ||
| 817 | * same counts (same shared data); the leader additionally sizes the | ||
| 818 | * subtree-ring slot from them. */ | ||
| 819 | extern void prism_pbuild_count_children( | ||
| 820 | PrismDsmSamples *dsm_samples, | ||
| 821 | PrismDsmRootAssign *dsm_ra, | ||
| 822 | int nparticipants, | ||
| 823 | uint32_t km_k, | ||
| 824 | uint32_t *out_counts); | ||
| 825 | |||
| 826 | /* | ||
| 827 | * The refine accumulator overlays the (dead) sample region: same base | ||
| 828 | * address, initialized by the leader after the last sample use and before | ||
| 829 | * the tree-ready barrier that workers pass ahead of the refine phase. | ||
| 830 | */ | ||
| 831 | static inline PrismDsmRefineAccum * | ||
| 832 | ✗ | prism_pbuild_refine_overlay(PrismDsmSamples *samples) | |
| 833 | { | ||
| 834 | ✗ | return (PrismDsmRefineAccum *)samples; | |
| 835 | } | ||
| 836 | |||
| 837 | extern bool prism_pbuild_launch( | ||
| 838 | struct ParallelContext *pcxt, | ||
| 839 | Barrier *barrier, | ||
| 840 | PrismBuildShared *shared); | ||
| 841 | |||
| 842 | /* ---------------------------------------------------------------- | ||
| 843 | * Leader-only subtree blob store — back-end-specific spillable storage | ||
| 844 | * | ||
| 845 | * The PLAN pass produces every subtree exactly once; the leader appends each | ||
| 846 | * blob here and, after computing the block layout, reads them back in the | ||
| 847 | * same order to stream centroid + head pages -- reading a blob back costs | ||
| 848 | * far less than re-running its clustering, and the streaming needs no | ||
| 849 | * worker participation. The PG implementation is a BufFile temp file: small | ||
| 850 | * blob sets stay | ||
| 851 | * in the kernel page cache, large ones spill to pgsql_tmp, so build memory | ||
| 852 | * stays bounded regardless of the partition count. Standalone keeps the | ||
| 853 | * blobs in memory (in-memory engine). Sequential put/rewind/get only. | ||
| 854 | * ---------------------------------------------------------------- */ | ||
| 855 | |||
| 856 | typedef struct PrismBlobStore PrismBlobStore; | ||
| 857 | |||
| 858 | extern PrismBlobStore *prism_pbuild_blobstore_begin(void); | ||
| 859 | extern void prism_pbuild_blobstore_put( | ||
| 860 | PrismBlobStore *bs, const void *blob, uint64_t size); | ||
| 861 | extern void prism_pbuild_blobstore_rewind(PrismBlobStore *bs); | ||
| 862 | /* Read the next blob into buf (capacity max_size); returns its size. */ | ||
| 863 | extern uint64_t | ||
| 864 | prism_pbuild_blobstore_get(PrismBlobStore *bs, void *buf, uint64_t max_size); | ||
| 865 | extern void prism_pbuild_blobstore_end(PrismBlobStore *bs); | ||
| 866 | |||
| 867 | /* ---------------------------------------------------------------- | ||
| 868 | * Build configuration — back-end-neutral input to the setup seam | ||
| 869 | * | ||
| 870 | * The fields the setup seam needs to size and populate the shared state. The | ||
| 871 | * PG path fills this from its opclass-resolved PrismBuildParams; the | ||
| 872 | * standalone path fills it from its own config. Keeping it neutral lets one | ||
| 873 | * setup-seam signature serve both back-ends. | ||
| 874 | * ---------------------------------------------------------------- */ | ||
| 875 | |||
| 876 | typedef struct PrismBuildConfig | ||
| 877 | { | ||
| 878 | Dimension dim; | ||
| 879 | DistanceMetric metric; | ||
| 880 | PrismCentroidFormat centroid_format; | ||
| 881 | uint32_t nlist; | ||
| 882 | uint32_t fan_out; | ||
| 883 | double soar_lambda; | ||
| 884 | double boundary_epsilon; | ||
| 885 | bool fastscan; | ||
| 886 | /* CREATE INDEX CONCURRENTLY: take weak locks + an MVCC scan snapshot. */ | ||
| 887 | bool concurrent; | ||
| 888 | } PrismBuildConfig; | ||
| 889 | |||
| 890 | /* ---------------------------------------------------------------- | ||
| 891 | * Leader setup seam — back-end-specific (parallel_backend.c for PG) | ||
| 892 | * | ||
| 893 | * Leader-side runtime state: the parallel context, the shared regions, and the | ||
| 894 | * derived sizes. prism_pbuild_setup_shared allocates and populates them (a DSM | ||
| 895 | * segment + regions in PG, a heap arena in standalone) and fills this struct; | ||
| 896 | * the rest of the driver consumes it. Returns false (after teardown) if the | ||
| 897 | * back-end could not start. The WalUsage/BufferUsage fields are PG | ||
| 898 | * instrumentation, unused in standalone. | ||
| 899 | * ---------------------------------------------------------------- */ | ||
| 900 | |||
| 901 | typedef struct PrismPBuildLeader | ||
| 902 | { | ||
| 903 | struct ParallelContext *pcxt; | ||
| 904 | PrismBuildShared *shared; | ||
| 905 | Barrier *barrier; | ||
| 906 | PrismDsmSamples *dsm_samples; | ||
| 907 | /* Back-end token for releasing the sample region early (PG: the DSM | ||
| 908 | * segment the samples live in; standalone: NULL). */ | ||
| 909 | void *sample_seg; | ||
| 910 | char *centroids_base; | ||
| 911 | float *cents; | ||
| 912 | char *km_workers_base; | ||
| 913 | PrismDsmRootAssign *dsm_ra; | ||
| 914 | struct WalUsage *walusage; | ||
| 915 | struct BufferUsage *bufferusage; | ||
| 916 | int nparticipants; | ||
| 917 | uint32_t km_k; | ||
| 918 | uint32_t max_per_worker; | ||
| 919 | Dimension dim; | ||
| 920 | uint32_t nlist; | ||
| 921 | uint64_t rabitq_seed; | ||
| 922 | uint32_t fan_out; | ||
| 923 | Size dsm_total; /* committed DSM chunk bytes (for the | ||
| 924 | planned-allocation introspection line) */ | ||
| 925 | } PrismPBuildLeader; | ||
| 926 | |||
| 927 | extern bool prism_pbuild_setup_shared( | ||
| 928 | PrismPBuildLeader *lead, | ||
| 929 | Relation heap, | ||
| 930 | Relation index, | ||
| 931 | const PrismBuildConfig *config, | ||
| 932 | int nworkers); | ||
| 933 | |||
| 934 | /* | ||
| 935 | * Re-initialize the scan for the posting pass (the sampling pass consumed the | ||
| 936 | * first). Back-end seam: resets the PG parallel-scan descriptor or, in | ||
| 937 | * standalone, the work-stealing cursor. | ||
| 938 | */ | ||
| 939 | extern void prism_pbuild_rescan(Relation heap, PrismBuildShared *shared); | ||
| 940 | |||
| 941 | /* ---------------------------------------------------------------- | ||
| 942 | * Posting sort seam — cluster-keyed external sort of encoded entries. | ||
| 943 | * | ||
| 944 | * Back-end seam (like prism_build_scan): the PG back-end implements it with a | ||
| 945 | * parallel tuplesort whose shared coordinator (Sharedsort) lives in the build | ||
| 946 | * DSM; the standalone back-end with per-participant in-memory arrays merged | ||
| 947 | * and sorted by cluster. The shared phase-3 code | ||
| 948 | * (parallel_build_worker/leader) calls only this seam, so the posting build | ||
| 949 | * stays bounded by maintenance_work_mem without the standalone core depending | ||
| 950 | * on PostgreSQL's tuplesort. | ||
| 951 | * | ||
| 952 | * Lifecycle: | ||
| 953 | * leader, in setup_shared before launch: size the DSM region with | ||
| 954 | * prism_pbuild_sort_shared_size(); after launch: | ||
| 955 | * prism_pbuild_sort_shared_init(). each worker: begin(is_leader=false) -> | ||
| 956 | * put... -> performsort -> end. leader: begin(is_leader=true) -> | ||
| 957 | * performsort -> getnext... -> end. Entries are fixed-size opaque blobs sorted | ||
| 958 | * by the uint32 cluster key. `seg` and `region` are void* to keep PostgreSQL | ||
| 959 | * types out of the shared header (the PG impl casts seg to dsm_segment*); | ||
| 960 | * `region` is the PRISM_DSM_KEY_SORTSHARED chunk. | ||
| 961 | * ---------------------------------------------------------------- */ | ||
| 962 | typedef struct PrismSorter PrismSorter; | ||
| 963 | |||
| 964 | extern Size prism_pbuild_sort_shared_size(int nparticipants); | ||
| 965 | extern void | ||
| 966 | prism_pbuild_sort_shared_init(void *region, int nparticipants, void *seg); | ||
| 967 | extern PrismSorter *prism_pbuild_sort_begin( | ||
| 968 | void *region, | ||
| 969 | void *seg, | ||
| 970 | int participant, | ||
| 971 | int nparticipants, | ||
| 972 | bool is_leader, | ||
| 973 | uint32_t entry_size, | ||
| 974 | int work_mem_kb); | ||
| 975 | extern void prism_pbuild_sort_put( | ||
| 976 | PrismSorter *sorter, uint32_t cluster, const void *entry); | ||
| 977 | extern void prism_pbuild_sort_performsort(PrismSorter *sorter); | ||
| 978 | extern bool prism_pbuild_sort_getnext( | ||
| 979 | PrismSorter *sorter, uint32_t *cluster, const void **entry); | ||
| 980 | extern void prism_pbuild_sort_end(PrismSorter *sorter); | ||
| 981 | |||
| 982 | /* | ||
| 983 | * Build every cluster's posting list from a populated (not yet performsorted) | ||
| 984 | * cluster-keyed sorter. Shared by the serial build and the parallel leader: | ||
| 985 | * performsort, then read entries grouped by cluster and write each list with a | ||
| 986 | * single resident page builder. Cluster c's head is the formula first_posting | ||
| 987 | * + c (a pre-extended head region); continuation pages are appended at the end | ||
| 988 | * of the relation and chained, so no O(nlist) reserve is needed. Ends the | ||
| 989 | * sorter. | ||
| 990 | */ | ||
| 991 | extern void prism_posting_build_lists( | ||
| 992 | PrismSorter *sorter, | ||
| 993 | VsStorage *storage, | ||
| 994 | uint32_t nlist, | ||
| 995 | Dimension dim, | ||
| 996 | bool fastscan, | ||
| 997 | const RaBitQParams *rq_params, | ||
| 998 | BlockNumber first_posting); | ||
| 999 | |||
| 1000 | /* ---------------------------------------------------------------- | ||
| 1001 | * Parallel build entry — shared driver (parallel_build_leader.c) | ||
| 1002 | * | ||
| 1003 | * Runs sampling + k-means + the bounded streaming posting build over the heap | ||
| 1004 | * (its vectors) into storage, returning the centroid tree and per-list posting | ||
| 1005 | * heads. Returns false if parallelism could not start, so the caller falls | ||
| 1006 | * back to a serial build. The PG caller fills config from its resolved build | ||
| 1007 | * params; the standalone caller fills it from its index config. | ||
| 1008 | * | ||
| 1009 | * prog is the build-progress reporting seam (phases + % to | ||
| 1010 | * pg_stat_progress_create_index, per-phase stats under prism.log_build_stats). | ||
| 1011 | * The PG caller passes its PrismBuildProgress so the parallel phases surface | ||
| 1012 | * the same way the serial ones do; the standalone caller passes NULL (the seam | ||
| 1013 | * is a no-op stub there). | ||
| 1014 | * ---------------------------------------------------------------- */ | ||
| 1015 | |||
| 1016 | struct PrismBuildProgress; /* index/build_progress.h */ | ||
| 1017 | |||
| 1018 | extern bool do_parallel_build( | ||
| 1019 | Relation heap, | ||
| 1020 | Relation index, | ||
| 1021 | struct IndexInfo *index_info, | ||
| 1022 | const PrismBuildConfig *config, | ||
| 1023 | VsStorage *storage, | ||
| 1024 | struct PrismBuildProgress *prog, | ||
| 1025 | /* No in-RAM tree is materialized; the streamed tree's shape (leaf | ||
| 1026 | * count + depth) comes back through these for the caller's metadata | ||
| 1027 | * write. *out_nlist == 0 means the heap had no indexable rows. */ | ||
| 1028 | uint32_t *out_nlist, | ||
| 1029 | uint8_t *out_tree_nlevels, | ||
| 1030 | double *out_heap_tuples, | ||
| 1031 | double *out_indtuples, | ||
| 1032 | double *out_soar_dupes, | ||
| 1033 | /* The build routes page-backed, so it writes the centroid pages + | ||
| 1034 | * heads into `storage` before the scan and computes the global mean; | ||
| 1035 | * *out_global_mean is the (owned) mean the centroid pages encode | ||
| 1036 | * against, so the caller's metadata write matches. May be NULL. */ | ||
| 1037 | float **out_global_mean, | ||
| 1038 | /* The posting-head base: cluster c's head is *out_first_posting + c. | ||
| 1039 | * Callers that need to locate head pages after the build (e.g. the | ||
| 1040 | * standalone driver + its tests) capture it; may be NULL. */ | ||
| 1041 | BlockNumber *out_first_posting); | ||
| 1042 | |||
| 1043 | #endif /* PRISM_PARALLEL_BUILD_H */ | ||
| 1044 |