summa_core/structures/vector/scann/
payload.rs

1use super::{ScannConfig, ScannEncoding, ScannFormatError, ScannResult, ScannTrainedArtifact};
2
3const MAGIC: &[u8; 8] = b"HSCNSEGM";
4/// Version 2: FastScan groups use the block-pair layout described in
5/// `docs/fast-scan-layout-v2.md`. Version 1 groups (one block per 16-byte
6/// word, two rows per byte) are refused; rebuild the generation.
7pub const SCANN_SEGMENT_PAYLOAD_VERSION: u16 = 2;
8
9/// One immutable physical leaf extent. Document IDs are little-endian u32
10/// values relative to `doc_base`; ordinals are little-endian u16 values.
11/// `codes` remain opaque to merging and are moved without re-encoding.
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct ScannLeafRun {
14    pub leaf_id: u32,
15    pub doc_base: u32,
16    pub row_count: u32,
17    pub doc_ids_le: Vec<u8>,
18    pub ordinals_le: Vec<u8>,
19    pub codes: Vec<u8>,
20}
21
22impl ScannLeafRun {
23    pub fn from_rows(
24        leaf_id: u32,
25        doc_base: u32,
26        doc_ids: &[u32],
27        ordinals: &[u16],
28        codes: Vec<u8>,
29        encoding: ScannEncoding,
30        dimension: u32,
31    ) -> ScannResult<Self> {
32        if doc_ids.len() != ordinals.len() {
33            return Err(ScannFormatError::new(
34                "ScaNN leaf document and ordinal columns have different lengths",
35            ));
36        }
37        let row_count = u32::try_from(doc_ids.len())
38            .map_err(|_| ScannFormatError::new("ScaNN leaf row count exceeds u32"))?;
39        let mut doc_ids_le = Vec::with_capacity(doc_ids.len().saturating_mul(4));
40        for &doc_id in doc_ids {
41            doc_ids_le.extend_from_slice(&doc_id.to_le_bytes());
42        }
43        let mut ordinals_le = Vec::with_capacity(ordinals.len().saturating_mul(2));
44        for &ordinal in ordinals {
45            ordinals_le.extend_from_slice(&ordinal.to_le_bytes());
46        }
47        let run = Self {
48            leaf_id,
49            doc_base,
50            row_count,
51            doc_ids_le,
52            ordinals_le,
53            codes,
54        };
55        run.validate(encoding, dimension, u32::MAX)?;
56        Ok(run)
57    }
58
59    fn validate(
60        &self,
61        encoding: ScannEncoding,
62        dimension: u32,
63        segment_docs: u32,
64    ) -> ScannResult<()> {
65        let rows = self.row_count as usize;
66        if self.doc_ids_le.len() != rows.saturating_mul(4)
67            || self.ordinals_le.len() != rows.saturating_mul(2)
68            || self.codes.len() != encoding.leaf_code_bytes(dimension, rows)?
69        {
70            return Err(ScannFormatError::new(
71                "ScaNN leaf run columns do not match its row count",
72            ));
73        }
74        for chunk in self.doc_ids_le.chunks_exact(4) {
75            let local = u32::from_le_bytes(chunk.try_into().unwrap());
76            let effective = self
77                .doc_base
78                .checked_add(local)
79                .ok_or_else(|| ScannFormatError::new("ScaNN document ID overflows u32"))?;
80            if effective >= segment_docs {
81                return Err(ScannFormatError::new(
82                    "ScaNN leaf run document ID is outside its segment",
83                ));
84            }
85        }
86        Ok(())
87    }
88}
89
90/// Segment-local leaf runs tied to one global trained artifact.
91#[derive(Debug, Clone, PartialEq, Eq)]
92pub struct ScannSegmentPayload {
93    pub artifact_id: u64,
94    pub generation: u64,
95    pub dimension: u32,
96    pub encoding: ScannEncoding,
97    pub num_leaves: u32,
98    pub doc_count: u32,
99    runs: Vec<ScannLeafRun>,
100}
101
102impl ScannSegmentPayload {
103    pub fn new(
104        artifact: &ScannTrainedArtifact,
105        doc_count: u32,
106        runs: Vec<ScannLeafRun>,
107    ) -> ScannResult<Self> {
108        Self::from_generation(
109            &artifact.config,
110            artifact.generation,
111            artifact.artifact_id,
112            doc_count,
113            runs,
114        )
115    }
116
117    /// Construct from the already-validated generation header without
118    /// decoding the global centroid plane into an owned artifact.
119    pub fn from_generation(
120        config: &ScannConfig,
121        generation: u64,
122        artifact_id: u64,
123        doc_count: u32,
124        mut runs: Vec<ScannLeafRun>,
125    ) -> ScannResult<Self> {
126        runs.sort_by_key(|run| run.leaf_id);
127        let payload = Self {
128            artifact_id,
129            generation,
130            dimension: config.dimension,
131            encoding: config.encoding,
132            num_leaves: config.num_leaves,
133            doc_count,
134            runs,
135        };
136        payload.validate()?;
137        Ok(payload)
138    }
139
140    pub fn runs(&self) -> &[ScannLeafRun] {
141        &self.runs
142    }
143
144    pub fn validate_against(&self, artifact: &ScannTrainedArtifact) -> ScannResult<()> {
145        artifact.validate()?;
146        if self.artifact_id != artifact.artifact_id
147            || self.generation != artifact.generation
148            || self.dimension != artifact.config.dimension
149            || self.encoding != artifact.config.encoding
150            || self.num_leaves != artifact.config.num_leaves
151        {
152            return Err(ScannFormatError::new(
153                "ScaNN segment payload does not match the global trained generation",
154            ));
155        }
156        self.validate()
157    }
158
159    /// Merge contiguous source segments by rebasing leaf-run document bases.
160    /// Code, ordinal, and relative-document buffers are moved as-is; no vector
161    /// is decoded, reassigned, or re-encoded and no codebook is retrained.
162    pub fn merge_contiguous(segments: impl IntoIterator<Item = Self>) -> ScannResult<Self> {
163        let mut segments = segments.into_iter();
164        let mut merged = segments
165            .next()
166            .ok_or_else(|| ScannFormatError::new("cannot merge zero ScaNN segments"))?;
167        merged.validate()?;
168        let mut next_doc_base = merged.doc_count;
169        for mut source in segments {
170            source.validate()?;
171            merged.ensure_compatible(&source)?;
172            for run in &mut source.runs {
173                run.doc_base = run
174                    .doc_base
175                    .checked_add(next_doc_base)
176                    .ok_or_else(|| ScannFormatError::new("merged ScaNN doc base exceeds u32"))?;
177            }
178            next_doc_base = next_doc_base
179                .checked_add(source.doc_count)
180                .ok_or_else(|| ScannFormatError::new("merged ScaNN document count exceeds u32"))?;
181            merged.runs.append(&mut source.runs);
182        }
183        merged.doc_count = next_doc_base;
184        // Stable leaf ordering keeps source segment order within a leaf while
185        // moving only run descriptors (the large column allocations stay put).
186        merged.runs.sort_by_key(|run| run.leaf_id);
187        merged.validate()?;
188        Ok(merged)
189    }
190
191    pub fn to_bytes(&self) -> ScannResult<Vec<u8>> {
192        self.validate()?;
193        let mut output = Vec::new();
194        output.extend_from_slice(MAGIC);
195        push_u16(&mut output, SCANN_SEGMENT_PAYLOAD_VERSION);
196        push_u16(&mut output, 0);
197        push_u64(&mut output, self.artifact_id);
198        push_u64(&mut output, self.generation);
199        push_u32(&mut output, self.dimension);
200        output.push(self.encoding.tag());
201        let (dimensions_per_block, bits_per_code) = self.encoding.parameters();
202        output.push(bits_per_code);
203        push_u16(&mut output, dimensions_per_block);
204        push_u32(&mut output, self.num_leaves);
205        push_u32(&mut output, self.doc_count);
206        push_u32(
207            &mut output,
208            u32::try_from(self.runs.len())
209                .map_err(|_| ScannFormatError::new("ScaNN run count exceeds u32"))?,
210        );
211        for run in &self.runs {
212            push_u32(&mut output, run.leaf_id);
213            push_u32(&mut output, run.doc_base);
214            push_u32(&mut output, run.row_count);
215            push_u64(&mut output, run.doc_ids_le.len() as u64);
216            push_u64(&mut output, run.ordinals_le.len() as u64);
217            push_u64(&mut output, run.codes.len() as u64);
218        }
219        for run in &self.runs {
220            output.extend_from_slice(&run.doc_ids_le);
221            output.extend_from_slice(&run.ordinals_le);
222            output.extend_from_slice(&run.codes);
223        }
224        Ok(output)
225    }
226
227    pub fn from_bytes(bytes: &[u8]) -> ScannResult<Self> {
228        let mut input = Input::new(bytes);
229        if input.take(8)? != MAGIC {
230            return Err(ScannFormatError::new("invalid ScaNN segment payload magic"));
231        }
232        let version = input.u16()?;
233        if version != SCANN_SEGMENT_PAYLOAD_VERSION {
234            return Err(ScannFormatError::new(format!(
235                "unsupported ScaNN segment payload version {version}; reader supports {SCANN_SEGMENT_PAYLOAD_VERSION}"
236            )));
237        }
238        if input.u16()? != 0 {
239            return Err(ScannFormatError::new(
240                "ScaNN segment reserved field is non-zero",
241            ));
242        }
243        let artifact_id = input.u64()?;
244        let generation = input.u64()?;
245        let dimension = input.u32()?;
246        let encoding_tag = input.u8()?;
247        let bits_per_code = input.u8()?;
248        let dimensions_per_block = input.u16()?;
249        let encoding =
250            ScannEncoding::from_parts(encoding_tag, dimensions_per_block, bits_per_code)?;
251        let num_leaves = input.u32()?;
252        let doc_count = input.u32()?;
253        let run_count = input.u32()? as usize;
254        if run_count > input.remaining() / 36 {
255            return Err(ScannFormatError::new(
256                "ScaNN segment run directory is truncated",
257            ));
258        }
259        let mut directory = Vec::with_capacity(run_count);
260        for _ in 0..run_count {
261            directory.push((
262                input.u32()?,
263                input.u32()?,
264                input.u32()?,
265                input.usize()?,
266                input.usize()?,
267                input.usize()?,
268            ));
269        }
270        let mut runs = Vec::with_capacity(run_count);
271        for (leaf_id, doc_base, row_count, docs_len, ordinals_len, codes_len) in directory {
272            runs.push(ScannLeafRun {
273                leaf_id,
274                doc_base,
275                row_count,
276                doc_ids_le: input.take(docs_len)?.to_vec(),
277                ordinals_le: input.take(ordinals_len)?.to_vec(),
278                codes: input.take(codes_len)?.to_vec(),
279            });
280        }
281        if !input.is_empty() {
282            return Err(ScannFormatError::new(
283                "ScaNN segment payload has trailing bytes",
284            ));
285        }
286        let payload = Self {
287            artifact_id,
288            generation,
289            dimension,
290            encoding,
291            num_leaves,
292            doc_count,
293            runs,
294        };
295        payload.validate()?;
296        Ok(payload)
297    }
298
299    fn ensure_compatible(&self, other: &Self) -> ScannResult<()> {
300        if self.artifact_id != other.artifact_id
301            || self.generation != other.generation
302            || self.dimension != other.dimension
303            || self.encoding != other.encoding
304            || self.num_leaves != other.num_leaves
305        {
306            return Err(ScannFormatError::new(
307                "cannot merge ScaNN segments from different trained generations",
308            ));
309        }
310        Ok(())
311    }
312
313    fn validate(&self) -> ScannResult<()> {
314        if self.artifact_id == 0 || self.generation == 0 || self.num_leaves == 0 {
315            return Err(ScannFormatError::new(
316                "invalid ScaNN segment compatibility metadata",
317            ));
318        }
319        self.encoding.row_code_bytes(self.dimension)?;
320        let mut previous_leaf = None;
321        for run in &self.runs {
322            if run.leaf_id >= self.num_leaves
323                || previous_leaf.is_some_and(|leaf| leaf > run.leaf_id)
324            {
325                return Err(ScannFormatError::new(
326                    "ScaNN leaf runs are out of range or unsorted",
327                ));
328            }
329            previous_leaf = Some(run.leaf_id);
330            run.validate(self.encoding, self.dimension, self.doc_count)?;
331        }
332        Ok(())
333    }
334}
335
336fn push_u16(output: &mut Vec<u8>, value: u16) {
337    output.extend_from_slice(&value.to_le_bytes());
338}
339
340fn push_u32(output: &mut Vec<u8>, value: u32) {
341    output.extend_from_slice(&value.to_le_bytes());
342}
343
344fn push_u64(output: &mut Vec<u8>, value: u64) {
345    output.extend_from_slice(&value.to_le_bytes());
346}
347
348struct Input<'a> {
349    bytes: &'a [u8],
350    offset: usize,
351}
352
353impl<'a> Input<'a> {
354    fn new(bytes: &'a [u8]) -> Self {
355        Self { bytes, offset: 0 }
356    }
357
358    fn take(&mut self, len: usize) -> ScannResult<&'a [u8]> {
359        let end = self
360            .offset
361            .checked_add(len)
362            .ok_or_else(|| ScannFormatError::new("ScaNN segment offset overflows"))?;
363        let value = self
364            .bytes
365            .get(self.offset..end)
366            .ok_or_else(|| ScannFormatError::new("truncated ScaNN segment payload"))?;
367        self.offset = end;
368        Ok(value)
369    }
370
371    fn u16(&mut self) -> ScannResult<u16> {
372        Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
373    }
374
375    fn u8(&mut self) -> ScannResult<u8> {
376        Ok(self.take(1)?[0])
377    }
378
379    fn u32(&mut self) -> ScannResult<u32> {
380        Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
381    }
382
383    fn u64(&mut self) -> ScannResult<u64> {
384        Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
385    }
386
387    fn usize(&mut self) -> ScannResult<usize> {
388        usize::try_from(self.u64()?)
389            .map_err(|_| ScannFormatError::new("ScaNN segment length exceeds usize"))
390    }
391
392    fn is_empty(&self) -> bool {
393        self.offset == self.bytes.len()
394    }
395
396    fn remaining(&self) -> usize {
397        self.bytes.len() - self.offset
398    }
399}
400
401#[cfg(test)]
402mod tests {
403    use super::super::{ScannConfig, ScannRoutingLevel, ScannTrainedArtifact};
404    use super::*;
405
406    fn artifact(generation: u64) -> ScannTrainedArtifact {
407        ScannTrainedArtifact::new(
408            generation,
409            100_000,
410            ScannConfig {
411                dimension: 16,
412                tree_levels: 1,
413                num_leaves: 2,
414                encoding: ScannEncoding::BinaryHamming,
415            },
416            vec![ScannRoutingLevel {
417                centroid_count: 2,
418                centroid_codes: vec![0, 0, 0xff, 0xff],
419                minimums: Vec::new(),
420                steps: Vec::new(),
421                child_offsets: Vec::new(),
422            }],
423            None,
424        )
425        .unwrap()
426    }
427
428    fn segment(artifact: &ScannTrainedArtifact, doc: u32, leaf: u32) -> ScannSegmentPayload {
429        let run = ScannLeafRun::from_rows(
430            leaf,
431            0,
432            &[doc],
433            &[0],
434            vec![doc as u8, 0],
435            ScannEncoding::BinaryHamming,
436            16,
437        )
438        .unwrap();
439        ScannSegmentPayload::new(artifact, doc + 1, vec![run]).unwrap()
440    }
441
442    #[test]
443    fn compatible_segment_merge_moves_code_buffers_and_rebases_only_run_metadata() {
444        let artifact = artifact(9);
445        let left = segment(&artifact, 0, 1);
446        let right = segment(&artifact, 0, 1);
447        let left_codes = left.runs()[0].codes.as_ptr();
448        let right_codes = right.runs()[0].codes.as_ptr();
449        let merged = ScannSegmentPayload::merge_contiguous([left, right]).unwrap();
450
451        assert_eq!(merged.doc_count, 2);
452        assert_eq!(merged.runs().len(), 2);
453        assert_eq!(merged.runs()[0].doc_base, 0);
454        assert_eq!(merged.runs()[1].doc_base, 1);
455        assert_eq!(merged.runs()[0].codes.as_ptr(), left_codes);
456        assert_eq!(merged.runs()[1].codes.as_ptr(), right_codes);
457    }
458
459    #[test]
460    fn segment_merge_refuses_different_global_generations() {
461        let first_artifact = artifact(10);
462        let second_artifact = artifact(11);
463        let first = segment(&first_artifact, 0, 0);
464        let second = segment(&second_artifact, 0, 0);
465        let error = ScannSegmentPayload::merge_contiguous([first, second]).unwrap_err();
466        assert!(error.to_string().contains("different trained generations"));
467    }
468
469    #[test]
470    fn segment_payload_round_trip_preserves_binary_leaf_runs() {
471        let artifact = artifact(12);
472        let segment = segment(&artifact, 0, 1);
473        let bytes = segment.to_bytes().unwrap();
474        let decoded = ScannSegmentPayload::from_bytes(&bytes).unwrap();
475        assert_eq!(decoded, segment);
476        decoded.validate_against(&artifact).unwrap();
477    }
478
479    #[test]
480    fn segment_payload_rejects_a_future_version() {
481        let artifact = artifact(13);
482        let segment = segment(&artifact, 0, 0);
483        let mut bytes = segment.to_bytes().unwrap();
484        bytes[8..10].copy_from_slice(&(SCANN_SEGMENT_PAYLOAD_VERSION + 1).to_le_bytes());
485        let error = ScannSegmentPayload::from_bytes(&bytes).unwrap_err();
486        assert!(error.to_string().contains("unsupported"));
487    }
488}
summa_core/structures/vector/scann/payload.rs

summa_core/structures/vector/scann/
payload.rs
