hermes_core/structures/vector/scann/
config.rs

1use super::{ScannFormatError, ScannResult};
2
3/// Maximum number of persisted routing centroid levels.
4///
5/// Matches the imported ScaNN artifact layout and AlloyDB-compatible knob.
6/// Three levels support billion-scale trees with configurable branching.
7pub const MAX_SCANN_TREE_LEVELS: u8 = 3;
8/// On-disk leaf identifiers are u32, but the imported routing format is
9/// intentionally capped to bound resident directories and route fan-out.
10pub const MAX_SCANN_LEAVES: u32 = 30_000_000;
11/// Below this many vectors, the fresh ScaNN implementation stays flat rather
12/// than training a partition tree.
13pub const MIN_POINTS_FOR_PARTITIONING: u64 = 100_000;
14/// Hardcoded minimum number of sampled rows per terminal leaf.
15///
16/// This is deliberately not a schema or server setting: a routing tree with
17/// fewer samples is not a viable trained geometry. Automatic geometry shrinks
18/// to fit the builder's sample budget; explicitly requested geometry is
19/// rejected when that budget cannot provide this floor.
20pub const MIN_PARTITION_TRAINING_POINTS_PER_LEAF: u64 = 8;
21/// Recall-oriented routing sample target per leaf centroid.
22pub const PARTITION_TRAINING_POINTS_PER_CENTROID: u64 = 200;
23/// Absolute routing/AH training sample target for small trained trees.
24pub const DEFAULT_TRAINING_SAMPLE_SIZE: u64 = 100_000;
25pub const SCANN_FAST_SCAN_LANES: usize = 32;
26
27/// Leaf representation used by the segment-local ScaNN payload.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum ScannEncoding {
30    /// Residual asymmetric-hash codes for floating-point embeddings.
31    AsymmetricHash {
32        dimensions_per_block: u16,
33        bits_per_code: u8,
34    },
35    /// Exact packed binary embeddings scored with Hamming distance.
36    BinaryHamming,
37}
38
39impl ScannEncoding {
40    pub(crate) fn tag(self) -> u8 {
41        match self {
42            Self::AsymmetricHash { .. } => 1,
43            Self::BinaryHamming => 2,
44        }
45    }
46
47    pub(crate) fn parameters(self) -> (u16, u8) {
48        match self {
49            Self::AsymmetricHash {
50                dimensions_per_block,
51                bits_per_code,
52            } => (dimensions_per_block, bits_per_code),
53            Self::BinaryHamming => (0, 0),
54        }
55    }
56
57    pub(crate) fn from_parts(tag: u8, dimensions_per_block: u16, bits: u8) -> ScannResult<Self> {
58        match (tag, dimensions_per_block, bits) {
59            (1, dimensions_per_block, bits_per_code) => Ok(Self::AsymmetricHash {
60                dimensions_per_block,
61                bits_per_code,
62            }),
63            (2, 0, 0) => Ok(Self::BinaryHamming),
64            _ => Err(ScannFormatError::new(
65                "invalid ScaNN leaf encoding or encoding parameters",
66            )),
67        }
68    }
69
70    pub fn row_code_bytes(self, dimension: u32) -> ScannResult<usize> {
71        let dimension = usize::try_from(dimension)
72            .map_err(|_| ScannFormatError::new("ScaNN dimension exceeds usize"))?;
73        match self {
74            Self::AsymmetricHash {
75                dimensions_per_block,
76                bits_per_code,
77            } => {
78                if dimensions_per_block == 0 || bits_per_code != 4 {
79                    return Err(ScannFormatError::new(
80                        "ScaNN AH encoding requires non-zero block dimensions and 4-bit codes",
81                    ));
82                }
83                let blocks = dimension.div_ceil(usize::from(dimensions_per_block));
84                blocks
85                    .checked_mul(usize::from(bits_per_code))
86                    .and_then(|bits| bits.checked_add(7))
87                    .map(|bits| bits / 8)
88                    .ok_or_else(|| ScannFormatError::new("ScaNN AH row size overflows usize"))
89            }
90            Self::BinaryHamming => {
91                if !dimension.is_multiple_of(8) {
92                    return Err(ScannFormatError::new(
93                        "binary ScaNN dimension must be a multiple of eight bits",
94                    ));
95                }
96                Ok(dimension / 8)
97            }
98        }
99    }
100
101    /// Encoded byte length for one leaf's corpus-sized code column. AH uses
102    /// the 32-lane FastScan v2 layout (two blocks per 32-byte word, odd block
103    /// counts padded; see `docs/fast-scan-layout-v2.md`) for complete groups
104    /// and compact row-major packing for the tail.
105    pub fn leaf_code_bytes(self, dimension: u32, rows: usize) -> ScannResult<usize> {
106        match self {
107            Self::BinaryHamming => self
108                .row_code_bytes(dimension)?
109                .checked_mul(rows)
110                .ok_or_else(|| ScannFormatError::new("binary ScaNN leaf size overflows")),
111            Self::AsymmetricHash {
112                dimensions_per_block,
113                ..
114            } => {
115                self.row_code_bytes(dimension)?;
116                let blocks = (dimension as usize).div_ceil(usize::from(dimensions_per_block));
117                let full_rows = rows / SCANN_FAST_SCAN_LANES;
118                let tail_rows = rows % SCANN_FAST_SCAN_LANES;
119                let full_block_bytes = super::packed_block_bytes(blocks)
120                    .ok_or_else(|| ScannFormatError::new("ScaNN FastScan block size overflows"))?;
121                let tail_row_bytes = blocks.div_ceil(2);
122                full_rows
123                    .checked_mul(full_block_bytes)
124                    .and_then(|bytes| {
125                        tail_rows
126                            .checked_mul(tail_row_bytes)
127                            .and_then(|tail| bytes.checked_add(tail))
128                    })
129                    .ok_or_else(|| ScannFormatError::new("ScaNN FastScan leaf size overflows"))
130            }
131        }
132    }
133}
134
135/// Index-scoped ScaNN training and layout configuration.
136#[derive(Debug, Clone, PartialEq, Eq)]
137pub struct ScannConfig {
138    /// Float dimensions for AH, bit dimensions for binary Hamming.
139    pub dimension: u32,
140    /// Configurable number of routing centroid levels.
141    pub tree_levels: u8,
142    /// Number of terminal leaves shared by every segment.
143    pub num_leaves: u32,
144    pub encoding: ScannEncoding,
145}
146
147#[derive(Debug, Clone, Copy, PartialEq, Eq)]
148pub enum ScannTrainingState {
149    AwaitingData { observed: u64, required: u64 },
150    Ready { observed: u64, required: u64 },
151}
152
153impl ScannConfig {
154    pub fn validate(&self) -> ScannResult<()> {
155        if self.dimension == 0 {
156            return Err(ScannFormatError::new(
157                "ScaNN vector dimension must be positive",
158            ));
159        }
160        if self.tree_levels == 0 || self.tree_levels > MAX_SCANN_TREE_LEVELS {
161            return Err(ScannFormatError::new(format!(
162                "ScaNN tree_levels must be in 1..={MAX_SCANN_TREE_LEVELS}"
163            )));
164        }
165        if self.num_leaves < 2 || self.num_leaves > MAX_SCANN_LEAVES {
166            return Err(ScannFormatError::new(
167                "ScaNN num_leaves must be in 2..=30,000,000",
168            ));
169        }
170        self.encoding.row_code_bytes(self.dimension)?;
171        Ok(())
172    }
173
174    /// Minimum viable sample. A trainer may reduce the desired sample under a
175    /// memory budget only while retaining at least this many rows.
176    pub fn minimum_training_sample(&self) -> ScannResult<u64> {
177        self.validate()?;
178        u64::from(self.num_leaves)
179            .checked_mul(MIN_PARTITION_TRAINING_POINTS_PER_LEAF)
180            .ok_or_else(|| ScannFormatError::new("ScaNN minimum training sample overflows u64"))
181    }
182
183    /// Recall-oriented sample target from the fresh ScaNN builder. The target
184    /// is capped by the observed corpus, but never silently changes geometry.
185    pub fn desired_training_sample(&self, observed_vectors: u64) -> ScannResult<u64> {
186        self.validate()?;
187        let desired = u64::from(self.num_leaves)
188            .checked_mul(PARTITION_TRAINING_POINTS_PER_CENTROID)
189            .ok_or_else(|| ScannFormatError::new("ScaNN training sample target overflows u64"))?
190            .max(DEFAULT_TRAINING_SAMPLE_SIZE);
191        Ok(observed_vectors.min(desired))
192    }
193
194    /// Minimum corpus size at which a requested partition geometry may train.
195    /// Before this threshold, serving must use the exact fallback.
196    pub fn effective_training_threshold(&self) -> ScannResult<u64> {
197        self.validate()?;
198        Ok(MIN_POINTS_FOR_PARTITIONING.max(self.minimum_training_sample()?))
199    }
200
201    pub fn training_state(&self, observed_vectors: u64) -> ScannResult<ScannTrainingState> {
202        let required = self.effective_training_threshold()?;
203        Ok(if observed_vectors < required {
204            ScannTrainingState::AwaitingData {
205                observed: observed_vectors,
206                required,
207            }
208        } else {
209            ScannTrainingState::Ready {
210                observed: observed_vectors,
211                required,
212            }
213        })
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220
221    fn config(encoding: ScannEncoding) -> ScannConfig {
222        ScannConfig {
223            dimension: 128,
224            tree_levels: 3,
225            num_leaves: 1_000_000,
226            encoding,
227        }
228    }
229
230    #[test]
231    fn scann_training_waits_for_partition_floor_and_minimum_leaf_sample() {
232        let config = config(ScannEncoding::BinaryHamming);
233        assert_eq!(config.effective_training_threshold().unwrap(), 8_000_000);
234        assert_eq!(
235            config.training_state(7_999_999).unwrap(),
236            ScannTrainingState::AwaitingData {
237                observed: 7_999_999,
238                required: 8_000_000,
239            }
240        );
241        assert!(matches!(
242            config.training_state(8_000_000).unwrap(),
243            ScannTrainingState::Ready { .. }
244        ));
245        assert_eq!(
246            config.desired_training_sample(8_000_000).unwrap(),
247            8_000_000
248        );
249
250        let mut smaller = config;
251        smaller.num_leaves = 1_000;
252        assert_eq!(smaller.effective_training_threshold().unwrap(), 100_000);
253        assert_eq!(smaller.desired_training_sample(1_000_000).unwrap(), 200_000);
254    }
255
256    #[test]
257    fn binary_scann_requires_a_byte_aligned_bit_dimension() {
258        let mut config = config(ScannEncoding::BinaryHamming);
259        config.dimension = 127;
260        assert!(config.validate().is_err());
261    }
262
263    #[test]
264    fn ah_row_size_is_derived_without_rounding_down() {
265        let encoding = ScannEncoding::AsymmetricHash {
266            dimensions_per_block: 2,
267            bits_per_code: 4,
268        };
269        assert_eq!(encoding.row_code_bytes(5).unwrap(), 2);
270        // FastScan v2 stores two blocks per 32-byte word: three blocks pad to
271        // two words (64 bytes) per complete 32-row group, and the 33rd row is
272        // a row-major tail of `ceil(3 / 2)` bytes.
273        assert_eq!(encoding.leaf_code_bytes(5, 32).unwrap(), 64);
274        assert_eq!(encoding.leaf_code_bytes(5, 33).unwrap(), 66);
275        // An even block count needs no padding.
276        assert_eq!(encoding.leaf_code_bytes(8, 32).unwrap(), 64);
277        assert_eq!(encoding.leaf_code_bytes(8, 64).unwrap(), 128);
278    }
279}
hermes_core/structures/vector/scann/config.rs

hermes_core/structures/vector/scann/
config.rs
