prolly/prolly/proximity/accelerator/
mod.rs

1mod r#async;
2pub mod catalog;
3pub mod composite;
4pub mod hnsw;
5pub mod pq;
6pub(crate) mod sq8;
7
8use self::composite::CompositeAccelerator;
9use self::hnsw::HnswIndex;
10use self::pq::ProductQuantizer;
11use super::ProximityTree;
12use crate::prolly::error::Error;
13use crate::prolly::store::Store;
14pub(crate) use r#async::AsyncCompositeBase;
15pub use r#async::{
16    AsyncAcceleratorCatalog, AsyncAcceleratorSet, AsyncCompositeAccelerator, AsyncHnswIndex,
17    AsyncProductQuantizer,
18};
19
20/// Validated source-bound derived accelerators available to one search.
21pub struct AcceleratorSet<S: Store> {
22    hnsw: Option<HnswIndex<S>>,
23    pq: Option<ProductQuantizer<S>>,
24    composite: Option<CompositeAccelerator<S>>,
25}
26
27impl<S: Store> Default for AcceleratorSet<S> {
28    fn default() -> Self {
29        Self {
30            hnsw: None,
31            pq: None,
32            composite: None,
33        }
34    }
35}
36
37impl<S: Store> AcceleratorSet<S> {
38    pub fn try_new(
39        source: &ProximityTree,
40        hnsw: Option<HnswIndex<S>>,
41        pq: Option<ProductQuantizer<S>>,
42    ) -> Result<Self, Error> {
43        if let Some(index) = &hnsw {
44            validate_binding(
45                source,
46                &index.source,
47                index.dimensions,
48                index.metric,
49                index.count,
50                "HNSW",
51            )?;
52        }
53        if let Some(index) = &pq {
54            validate_binding(
55                source,
56                &index.source,
57                index.dimensions,
58                index.metric,
59                index.count,
60                "product quantization",
61            )?;
62        }
63        Ok(Self {
64            hnsw,
65            pq,
66            composite: None,
67        })
68    }
69
70    pub fn empty() -> Self {
71        Self::default()
72    }
73
74    pub fn with_hnsw(mut self, source: &ProximityTree, index: HnswIndex<S>) -> Result<Self, Error> {
75        if self.hnsw.is_some() {
76            return Err(invalid("duplicate HNSW accelerator"));
77        }
78        validate_binding(
79            source,
80            &index.source,
81            index.dimensions,
82            index.metric,
83            index.count,
84            "HNSW",
85        )?;
86        self.hnsw = Some(index);
87        Ok(self)
88    }
89
90    pub fn with_pq(
91        mut self,
92        source: &ProximityTree,
93        index: ProductQuantizer<S>,
94    ) -> Result<Self, Error> {
95        if self.pq.is_some() {
96            return Err(invalid("duplicate product-quantization accelerator"));
97        }
98        validate_binding(
99            source,
100            &index.source,
101            index.dimensions,
102            index.metric,
103            index.count,
104            "product quantization",
105        )?;
106        self.pq = Some(index);
107        Ok(self)
108    }
109
110    pub fn with_composite(
111        mut self,
112        source: &ProximityTree,
113        accelerator: CompositeAccelerator<S>,
114    ) -> Result<Self, Error> {
115        if self.composite.is_some() {
116            return Err(invalid("duplicate composite accelerator"));
117        }
118        if accelerator.current_source != source.descriptor
119            || accelerator.dimensions != source.config.dimensions
120            || accelerator.metric != source.config.metric
121            || accelerator.current_count != source.count
122        {
123            return Err(invalid(
124                "composite accelerator is bound to a different source snapshot",
125            ));
126        }
127        self.composite = Some(accelerator);
128        Ok(self)
129    }
130
131    pub(crate) fn hnsw(&self) -> Option<&HnswIndex<S>> {
132        self.hnsw.as_ref()
133    }
134
135    pub(crate) fn pq(&self) -> Option<&ProductQuantizer<S>> {
136        self.pq.as_ref()
137    }
138
139    pub(crate) fn composite(&self) -> Option<&CompositeAccelerator<S>> {
140        self.composite.as_ref()
141    }
142}
143
144fn validate_binding(
145    source: &ProximityTree,
146    descriptor: &crate::prolly::cid::Cid,
147    dimensions: u32,
148    metric: super::DistanceMetric,
149    count: u64,
150    kind: &'static str,
151) -> Result<(), Error> {
152    if descriptor != &source.descriptor
153        || dimensions != source.config.dimensions
154        || metric != source.config.metric
155        || count != source.count
156    {
157        return Err(invalid(format!(
158            "{kind} accelerator is bound to a different source snapshot"
159        )));
160    }
161    Ok(())
162}
163
164fn invalid(reason: impl Into<String>) -> Error {
165    Error::InvalidProximitySearch {
166        reason: reason.into(),
167    }
168}
prolly/prolly/proximity/accelerator/mod.rs

prolly/prolly/proximity/accelerator/
mod.rs
