1use super::{ApproximatePreference, EligibilityCardinality, PreparedFilter, SearchRequest};
2use crate::prolly::error::Error;
3use crate::prolly::proximity::accelerator::AcceleratorSet;
4use crate::prolly::proximity::{CompositeAcceleratorConfig, CompositeBaseKind};
5use crate::prolly::proximity::{HnswConfig, ProductQuantizationConfig};
6use crate::prolly::proximity::{ProximityTree, SearchBackend, SearchPolicy};
7use crate::prolly::store::Store;
8
9pub const SEARCH_PLAN_FORMAT_VERSION: u8 = 3;
10
11#[derive(Clone, Debug, PartialEq, Eq)]
12pub enum SearchPlan {
13 Native,
14 EligibleExact {
15 key_count: u64,
16 source_bound: bool,
17 },
18 ProductQuantized {
19 rerank_target: usize,
20 direct_lookup: bool,
21 },
22 Hnsw {
23 ef_search: u32,
24 expansion_target: usize,
25 rerank_target: usize,
26 },
27 Composite {
28 base: Box<SearchPlan>,
29 delta_records: usize,
30 shadow_records: usize,
31 merge_target: usize,
32 },
33}
34
35#[derive(Clone, Debug, PartialEq, Eq)]
36pub struct SearchPlanSummary {
37 pub format_version: u8,
38 pub backend: SearchBackend,
39 pub eligible_exact_records: Option<u64>,
40 pub hnsw_ef_search: Option<u32>,
41 pub expansion_target: Option<usize>,
42 pub rerank_target: Option<usize>,
43 pub direct_lookup: bool,
44 pub composite_base_backend: Option<SearchBackend>,
45 pub composite_base: Option<Box<SearchPlanSummary>>,
46 pub delta_records: Option<usize>,
47 pub shadow_records: Option<usize>,
48}
49
50impl SearchPlan {
51 pub fn summary(&self) -> SearchPlanSummary {
52 let mut summary = SearchPlanSummary {
53 format_version: SEARCH_PLAN_FORMAT_VERSION,
54 backend: SearchBackend::Native,
55 eligible_exact_records: None,
56 hnsw_ef_search: None,
57 expansion_target: None,
58 rerank_target: None,
59 direct_lookup: false,
60 composite_base_backend: None,
61 composite_base: None,
62 delta_records: None,
63 shadow_records: None,
64 };
65 match self {
66 Self::Native => {}
67 Self::EligibleExact { key_count, .. } => {
68 summary.eligible_exact_records = Some(*key_count);
69 }
70 Self::ProductQuantized {
71 rerank_target,
72 direct_lookup,
73 } => {
74 summary.backend = SearchBackend::ProductQuantized;
75 summary.rerank_target = Some(*rerank_target);
76 summary.direct_lookup = *direct_lookup;
77 }
78 Self::Hnsw {
79 ef_search,
80 expansion_target,
81 rerank_target,
82 } => {
83 summary.backend = SearchBackend::Hnsw;
84 summary.hnsw_ef_search = Some(*ef_search);
85 summary.expansion_target = Some(*expansion_target);
86 summary.rerank_target = Some(*rerank_target);
87 }
88 Self::Composite {
89 base,
90 delta_records,
91 shadow_records,
92 merge_target,
93 } => {
94 summary.backend = SearchBackend::Composite;
95 summary.composite_base_backend = Some(base.summary().backend);
96 summary.composite_base = Some(Box::new(base.summary()));
97 summary.delta_records = Some(*delta_records);
98 summary.shadow_records = Some(*shadow_records);
99 summary.rerank_target = Some(*merge_target);
100 }
101 }
102 summary
103 }
104}
105
106pub(crate) struct CompositePlanInput<'a> {
107 pub base_kind: CompositeBaseKind,
108 pub hnsw: Option<&'a HnswConfig>,
109 pub pq: Option<&'a ProductQuantizationConfig>,
110 pub base_count: u64,
111 pub delta_count: u64,
112 pub shadow_count: u64,
113 pub config: &'a CompositeAcceleratorConfig,
114}
115
116pub(crate) fn plan_search<S>(
117 tree: &ProximityTree,
118 accelerators: &AcceleratorSet<S>,
119 request: &SearchRequest<'_>,
120 eligibility: &PreparedFilter<'_>,
121) -> Result<SearchPlan, Error>
122where
123 S: Store + Clone + Send + Sync,
124 S::Error: Send + Sync,
125{
126 plan_search_capabilities(
127 tree,
128 accelerators.hnsw().map(|index| index.config()),
129 accelerators.pq().map(|index| index.config()),
130 accelerators
131 .composite()
132 .map(|composite| CompositePlanInput {
133 base_kind: composite.base_kind(),
134 hnsw: composite.base.hnsw().map(|index| index.config()),
135 pq: composite.base.pq().map(|index| index.config()),
136 base_count: composite.base_count,
137 delta_count: composite.delta_count,
138 shadow_count: composite.shadow_count,
139 config: &composite.config,
140 }),
141 request,
142 eligibility,
143 )
144}
145
146pub(crate) fn plan_search_capabilities(
147 tree: &ProximityTree,
148 hnsw: Option<&HnswConfig>,
149 pq: Option<&ProductQuantizationConfig>,
150 composite: Option<CompositePlanInput<'_>>,
151 request: &SearchRequest<'_>,
152 eligibility: &PreparedFilter<'_>,
153) -> Result<SearchPlan, Error> {
154 let cardinality = eligibility.cardinality(tree.count);
155 match request.options.backend {
156 SearchBackend::Native => return Ok(SearchPlan::Native),
157 SearchBackend::Hnsw => {
158 ensure_approximate(request, "HNSW")?;
159 let config = hnsw.ok_or_else(|| invalid("forced HNSW backend is unavailable"))?;
160 return hnsw_plan(tree, config, request, cardinality);
161 }
162 SearchBackend::ProductQuantized => {
163 ensure_approximate(request, "product quantization")?;
164 let config =
165 pq.ok_or_else(|| invalid("forced product-quantized backend is unavailable"))?;
166 return pq_plan(tree, config, request, eligibility, cardinality);
167 }
168 SearchBackend::Composite => {
169 ensure_approximate(request, "composite")?;
170 return composite_plan(
171 tree,
172 composite.ok_or_else(|| invalid("forced composite backend is unavailable"))?,
173 request,
174 eligibility,
175 cardinality,
176 );
177 }
178 SearchBackend::Auto => {}
179 }
180
181 if let Some(plan) = eligible_exact_plan(tree, request, eligibility, cardinality)? {
182 if matches!(plan, SearchPlan::EligibleExact { key_count: 0, .. })
183 || request.policy == SearchPolicy::Exact
184 || request.options.planner.allow_exact_for_approximate
185 {
186 return Ok(plan);
187 }
188 }
189 if request.policy == SearchPolicy::Exact {
190 return Ok(SearchPlan::Native);
191 }
192
193 let preferences = match request.options.planner.approximate_preference {
194 ApproximatePreference::HnswFirst => [SearchBackend::Hnsw, SearchBackend::ProductQuantized],
195 ApproximatePreference::ProductQuantizedFirst => {
196 [SearchBackend::ProductQuantized, SearchBackend::Hnsw]
197 }
198 };
199 for backend in preferences {
200 let plan = match backend {
201 SearchBackend::Hnsw => hnsw
202 .map(|config| hnsw_plan(tree, config, request, cardinality))
203 .transpose()?,
204 SearchBackend::ProductQuantized => pq
205 .map(|config| pq_plan(tree, config, request, eligibility, cardinality))
206 .transpose()?,
207 SearchBackend::Native | SearchBackend::Composite | SearchBackend::Auto => None,
208 };
209 if plan
210 .as_ref()
211 .is_some_and(|plan| budget_admissible(plan, request))
212 {
213 return Ok(plan.expect("checked plan"));
214 }
215 }
216 if let Some(composite) = composite {
217 let plan = composite_plan(tree, composite, request, eligibility, cardinality)?;
218 if budget_admissible(&plan, request) {
219 return Ok(plan);
220 }
221 }
222 Ok(SearchPlan::Native)
223}
224
225fn composite_plan(
226 tree: &ProximityTree,
227 input: CompositePlanInput<'_>,
228 request: &SearchRequest<'_>,
229 eligibility: &PreparedFilter<'_>,
230 cardinality: EligibilityCardinality,
231) -> Result<SearchPlan, Error> {
232 let mut base = match input.base_kind {
233 CompositeBaseKind::Hnsw => hnsw_plan_for_count(
234 input.base_count,
235 input
236 .hnsw
237 .ok_or_else(|| invalid("composite HNSW configuration is absent"))?,
238 request,
239 cardinality,
240 )?,
241 CompositeBaseKind::ProductQuantized => pq_plan_for_count(
242 input.base_count,
243 input
244 .pq
245 .ok_or_else(|| invalid("composite PQ configuration is absent"))?,
246 request,
247 eligibility,
248 cardinality,
249 )?,
250 };
251 inflate_composite_base(
252 &mut base,
253 input.base_count,
254 input.shadow_count,
255 input.config.base_overfetch_multiplier,
256 )?;
257 let delta_records = usize::try_from(input.delta_count).unwrap_or(usize::MAX);
258 Ok(SearchPlan::Composite {
259 base: Box::new(base),
260 delta_records,
261 shadow_records: usize::try_from(input.shadow_count).unwrap_or(usize::MAX),
262 merge_target: request.k.min(tree.count as usize),
263 })
264}
265
266fn inflate_composite_base(
267 plan: &mut SearchPlan,
268 base_count: u64,
269 shadow_count: u64,
270 multiplier: u32,
271) -> Result<(), Error> {
272 let surviving = base_count.saturating_sub(shadow_count).max(1);
273 let inflate = |value: usize| -> Result<usize, Error> {
274 let scaled = (value as u128)
275 .checked_mul(u128::from(base_count))
276 .and_then(|value| value.checked_mul(u128::from(multiplier)))
277 .ok_or_else(|| invalid("composite base inflation overflow"))?
278 .div_ceil(u128::from(surviving));
279 Ok(usize::try_from(scaled.min(u128::from(base_count))).unwrap_or(usize::MAX))
280 };
281 match plan {
282 SearchPlan::Hnsw {
283 expansion_target,
284 rerank_target,
285 ..
286 } => {
287 *expansion_target = inflate(*expansion_target)?;
288 *rerank_target = inflate(*rerank_target)?;
289 }
290 SearchPlan::ProductQuantized {
291 rerank_target,
292 direct_lookup,
293 } => {
294 *rerank_target = inflate(*rerank_target)?;
295 *direct_lookup = false;
296 }
297 _ => return Err(invalid("composite base plan is not approximate")),
298 }
299 Ok(())
300}
301
302fn eligible_exact_plan(
303 tree: &ProximityTree,
304 request: &SearchRequest<'_>,
305 eligibility: &PreparedFilter<'_>,
306 cardinality: EligibilityCardinality,
307) -> Result<Option<SearchPlan>, Error> {
308 let EligibilityCardinality::Known(eligible) = cardinality else {
309 return Ok(None);
310 };
311 let Some((_, source_bound)) = eligibility.sorted_keys() else {
312 return Ok(None);
313 };
314 if eligible == 0 {
315 return Ok(Some(SearchPlan::EligibleExact {
316 key_count: 0,
317 source_bound,
318 }));
319 }
320 let ratio_numerator = u128::from(tree.count)
321 .checked_mul(u128::from(request.options.planner.eligible_exact_ratio_ppm))
322 .ok_or_else(|| invalid("eligible exact ratio overflow"))?;
323 let ratio_limit = ratio_numerator
324 .checked_add(999_999)
325 .ok_or_else(|| invalid("eligible exact ratio overflow"))?
326 / 1_000_000;
327 let ratio_limit = usize::try_from(ratio_limit).unwrap_or(usize::MAX);
328 let threshold = request.k.max(
329 request
330 .options
331 .planner
332 .eligible_exact_max_records
333 .min(ratio_limit),
334 );
335 Ok(
336 (eligible <= threshold as u64).then_some(SearchPlan::EligibleExact {
337 key_count: eligible,
338 source_bound,
339 }),
340 )
341}
342
343fn hnsw_plan(
344 tree: &ProximityTree,
345 config: &HnswConfig,
346 request: &SearchRequest<'_>,
347 cardinality: EligibilityCardinality,
348) -> Result<SearchPlan, Error> {
349 hnsw_plan_for_count(tree.count, config, request, cardinality)
350}
351
352fn hnsw_plan_for_count(
353 total_count: u64,
354 config: &HnswConfig,
355 request: &SearchRequest<'_>,
356 cardinality: EligibilityCardinality,
357) -> Result<SearchPlan, Error> {
358 let ef_search = request.options.hnsw.ef_search.unwrap_or(config.ef_search);
359 let base = usize::try_from(ef_search).unwrap_or(usize::MAX).max(
360 request
361 .k
362 .checked_mul(config.overfetch_multiplier as usize)
363 .ok_or_else(|| invalid("HNSW expansion target overflow"))?,
364 );
365 let expansion_target = match cardinality {
366 EligibilityCardinality::Known(0) => 0,
367 EligibilityCardinality::Known(eligible) => {
368 let numerator = (base as u128)
369 .checked_mul(u128::from(total_count))
370 .ok_or_else(|| invalid("HNSW selective expansion overflow"))?;
371 let target = numerator
372 .checked_add(u128::from(eligible) - 1)
373 .ok_or_else(|| invalid("HNSW selective expansion overflow"))?
374 / u128::from(eligible);
375 usize::try_from(target.min(u128::from(total_count))).unwrap_or(usize::MAX)
376 }
377 EligibilityCardinality::Unknown => base.min(total_count as usize),
378 };
379 let known_limit = match cardinality {
380 EligibilityCardinality::Known(count) => count,
381 EligibilityCardinality::Unknown => total_count,
382 };
383 let rerank_target = request
384 .k
385 .checked_mul(config.overfetch_multiplier as usize)
386 .ok_or_else(|| invalid("HNSW rerank target overflow"))?
387 .max(request.k)
388 .min(known_limit as usize)
389 .min(total_count as usize);
390 Ok(SearchPlan::Hnsw {
391 ef_search,
392 expansion_target,
393 rerank_target,
394 })
395}
396
397fn pq_plan(
398 tree: &ProximityTree,
399 config: &ProductQuantizationConfig,
400 request: &SearchRequest<'_>,
401 eligibility: &PreparedFilter<'_>,
402 cardinality: EligibilityCardinality,
403) -> Result<SearchPlan, Error> {
404 pq_plan_for_count(tree.count, config, request, eligibility, cardinality)
405}
406
407fn pq_plan_for_count(
408 total_count: u64,
409 config: &ProductQuantizationConfig,
410 request: &SearchRequest<'_>,
411 eligibility: &PreparedFilter<'_>,
412 cardinality: EligibilityCardinality,
413) -> Result<SearchPlan, Error> {
414 let multiplier = request
415 .options
416 .pq
417 .rerank_multiplier
418 .map(usize::from)
419 .unwrap_or(config.rerank_multiplier as usize);
420 let known_limit = match cardinality {
421 EligibilityCardinality::Known(count) => count,
422 EligibilityCardinality::Unknown => total_count,
423 };
424 let rerank_target = request
425 .k
426 .checked_mul(multiplier)
427 .ok_or_else(|| invalid("PQ rerank target overflow"))?
428 .max(request.k)
429 .min(known_limit as usize)
430 .min(total_count as usize);
431 let direct_lookup = eligibility.sorted_keys().is_some()
432 && known_limit <= request.options.planner.eligible_exact_max_records as u64;
433 Ok(SearchPlan::ProductQuantized {
434 rerank_target,
435 direct_lookup,
436 })
437}
438
439fn budget_admissible(plan: &SearchPlan, request: &SearchRequest<'_>) -> bool {
440 match plan {
441 SearchPlan::Hnsw {
442 expansion_target,
443 rerank_target,
444 ..
445 } => {
446 request
447 .budget
448 .max_nodes
449 .is_none_or(|limit| *expansion_target <= limit)
450 && request
451 .budget
452 .max_distance_evaluations
453 .is_none_or(|limit| expansion_target.saturating_add(*rerank_target) <= limit)
454 && request
455 .budget
456 .max_frontier_entries
457 .is_none_or(|limit| request.k <= limit)
458 }
459 SearchPlan::ProductQuantized { rerank_target, .. } => request
460 .budget
461 .max_distance_evaluations
462 .is_none_or(|limit| *rerank_target <= limit),
463 SearchPlan::Composite {
464 base,
465 delta_records,
466 shadow_records,
467 ..
468 } => {
469 let node_work = estimated_node_work(base)
470 .saturating_add(*shadow_records)
471 .saturating_add(delta_records.saturating_mul(2));
472 let distance_work = estimated_distance_work(base).saturating_add(*delta_records);
473 budget_admissible(base, request)
474 && request
475 .budget
476 .max_nodes
477 .is_none_or(|limit| node_work <= limit)
478 && request
479 .budget
480 .max_distance_evaluations
481 .is_none_or(|limit| distance_work <= limit)
482 }
483 SearchPlan::Native | SearchPlan::EligibleExact { .. } => true,
484 }
485}
486
487fn estimated_node_work(plan: &SearchPlan) -> usize {
488 match plan {
489 SearchPlan::Hnsw {
490 expansion_target,
491 rerank_target,
492 ..
493 } => expansion_target.saturating_add(*rerank_target),
494 SearchPlan::ProductQuantized { rerank_target, .. } => *rerank_target,
495 SearchPlan::Composite {
496 base,
497 delta_records,
498 shadow_records,
499 ..
500 } => estimated_node_work(base)
501 .saturating_add(*shadow_records)
502 .saturating_add(delta_records.saturating_mul(2)),
503 SearchPlan::Native | SearchPlan::EligibleExact { .. } => 0,
504 }
505}
506
507fn estimated_distance_work(plan: &SearchPlan) -> usize {
508 match plan {
509 SearchPlan::Hnsw {
510 expansion_target,
511 rerank_target,
512 ..
513 } => expansion_target.saturating_add(*rerank_target),
514 SearchPlan::ProductQuantized { rerank_target, .. } => *rerank_target,
515 SearchPlan::Composite {
516 base,
517 delta_records,
518 ..
519 } => estimated_distance_work(base).saturating_add(*delta_records),
520 SearchPlan::Native | SearchPlan::EligibleExact { .. } => 0,
521 }
522}
523
524fn ensure_approximate(request: &SearchRequest<'_>, backend: &str) -> Result<(), Error> {
525 if request.policy == SearchPolicy::Exact {
526 Err(invalid(format!("{backend} cannot satisfy exact search")))
527 } else {
528 Ok(())
529 }
530}
531
532fn invalid(reason: impl Into<String>) -> Error {
533 Error::InvalidProximitySearch {
534 reason: reason.into(),
535 }
536}