uqa_planner/cost_model/
vector.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! DiskANN work estimates from retained physical facts. These are uncalibrated work units, not elapsed time or measured SSD operations.
8
9use uqa_core::{DiskANNIndexStats, DiskANNQueryStats, VectorQueryRoute};
10
11#[derive(Debug, Clone, Copy, Default, PartialEq)]
12pub struct DiskANNWorkEstimate {
13    pub approximate_nodes: f64,
14    pub completion_nodes: f64,
15    pub pq_lookup_coordinates: f64,
16    pub pq_distance_evaluations: f64,
17    pub logical_page_requests: f64,
18    pub logical_page_bytes: f64,
19    pub page_rounds: f64,
20    pub side_vectors: f64,
21    pub changed_vectors: f64,
22    pub rerank_vectors: f64,
23    pub exact_vectors: f64,
24    pub resident_pq_payload_bytes: f64,
25    pub page_budget_sufficient: bool,
26    pub cpu: f64,
27    pub io: f64,
28}
29
30impl DiskANNWorkEstimate {
31    pub fn total(self) -> f64 {
32        self.cpu + self.io
33    }
34}
35
36/// None selects exact threshold retrieval; Some(k) selects a document candidate pool. Unknown current/change counts use the table population and stored tensor density as explicit estimates, never an invented measured zero.
37pub fn estimate_diskann(
38    stats: &DiskANNQueryStats,
39    k: Option<usize>,
40    table_documents: u64,
41) -> DiskANNWorkEstimate {
42    let query_route = stats.query_route;
43    let stats = &stats.index;
44    let mut work = DiskANNWorkEstimate::default();
45    let populations = stats.populations;
46    let dimensions = f64::from(stats.dimensions);
47    let density = populations
48        .base_documents
49        .filter(|count| *count != 0)
50        .map_or(1.0, |documents| {
51            (populations.base_vectors as f64 / documents as f64).max(1.0)
52        });
53    let current = populations
54        .current_vectors
55        .map_or(table_documents as f64 * density, |count| count as f64);
56    let centroids = stats
57        .pq_centroids
58        .map_or(populations.graph_nodes.min(256) as f64, f64::from);
59    work.resident_pq_payload_bytes =
60        dimensions * centroids * 8.0 + populations.graph_nodes as f64 * stats.pq_bytes as f64;
61    work.page_budget_sufficient = true;
62    if k == Some(0) {
63        return work;
64    }
65    if k.is_none() || query_route != VectorQueryRoute::Approximate {
66        work.exact_vectors = current;
67        work.cpu = dimensions * current;
68        return work;
69    }
70    let nodes = populations.graph_nodes as f64;
71    work.approximate_nodes = nodes.min(stats.search_list_size as f64);
72    let requested_vectors = k.unwrap_or_default() as f64 * density;
73    work.completion_nodes = (requested_vectors - work.approximate_nodes)
74        .max(0.0)
75        .min(nodes - work.approximate_nodes);
76    let expanded = work.approximate_nodes + work.completion_nodes;
77    work.pq_lookup_coordinates = dimensions * centroids;
78    work.pq_distance_evaluations = if nodes == 0.0 {
79        0.0
80    } else {
81        let degree = stats
82            .graph_edges
83            .map_or((stats.max_degree as f64).min(nodes - 1.0), |edges| {
84                edges as f64 / nodes
85            });
86        1.0 + work.approximate_nodes * degree
87    };
88    work.side_vectors = populations.side_vectors as f64;
89    work.changed_vectors = populations
90        .changed_vectors
91        .map_or(current, |count| count as f64);
92    // Each selected document is scored once with its complete tensor. Stored density is a planning assumption, not a bound on a particular tensor.
93    work.rerank_vectors = (expanded * density).min(current);
94    work.logical_page_requests = expanded * f64::from(stats.node_fragments);
95    work.logical_page_bytes = work.logical_page_requests * stats.page_bytes as f64;
96    work.page_rounds = page_rounds(stats, expanded);
97    work.page_budget_sufficient =
98        expanded == 0.0 || stats.reads.max_in_flight_page_bytes >= stats.page_bytes;
99    work.cpu = work.pq_lookup_coordinates
100        + work.pq_distance_evaluations * stats.pq_bytes as f64
101        + dimensions * (work.side_vectors * density + work.changed_vectors + work.rerank_vectors);
102    // Transfer and latency terms use fixed relative weights; no machine calibration is implied.
103    work.io = 8.0 * (work.logical_page_requests + work.page_rounds);
104    work
105}
106
107fn page_rounds(stats: &DiskANNIndexStats, expanded: f64) -> f64 {
108    if expanded == 0.0 {
109        return 0.0;
110    }
111    let beam = stats.beam_width.max(1) as f64;
112    let batch = stats
113        .reads
114        .max_batch_pages
115        .min(stats.reads.max_in_flight_page_bytes / stats.page_bytes.max(1));
116    // An inadmissible page budget remains visible in the estimate. Serial work describes what is required, not permission to dispatch with that allowance.
117    let batch = batch.max(1) as f64;
118    let concurrency = (stats.reads.read_concurrency.max(1) as f64).min(batch);
119    let full_beams = (expanded / beam).floor();
120    let remainder = expanded - full_beams * beam;
121    let fragments = f64::from(stats.node_fragments);
122    let rounds = |pages: f64| {
123        let complete = (pages / batch).floor();
124        complete * (batch / concurrency).ceil() + ((pages - complete * batch) / concurrency).ceil()
125    };
126    full_beams * rounds(beam * fragments) + rounds(remainder * fragments)
127}
128
129#[cfg(test)]
130mod tests;
uqa_planner/cost_model/vector.rs

uqa_planner/cost_model/
vector.rs
