velesdb_core/quantization/
pq_persistence.rs

1//! Persistence layer for Product Quantization codebooks and rotation matrices.
2//!
3//! Provides atomic file I/O using postcard serialization with crash-safe
4//! write-then-rename semantics. Extracted from [`super::pq`] to isolate
5//! the storage concern from the core PQ algorithm.
6
7use crate::error::Error;
8use serde::{Deserialize, Serialize};
9
10use super::pq::ProductQuantizer;
11
12/// Maximum accepted size of a persisted PQ artifact (codebook or rotation).
13///
14/// A valid codebook is `num_subspaces * num_centroids * subspace_dim` f32s
15/// plus small metadata; with the trained bounds (`num_subspaces <= 64`,
16/// `num_centroids <= u16::MAX`, `subspace_dim` modest) this stays well under
17/// 256 MiB. The cap rejects absurd/hostile files before they are decoded into
18/// a multi-gigabyte allocation. (Addresses the alloc-cap concern of #897.)
19const MAX_PQ_ARTIFACT_BYTES: u64 = 256 * 1024 * 1024;
20
21/// RF-2: Serializes `value` with postcard and atomically writes to `dir/filename`.
22///
23/// Write goes to `.tmp` suffix first, then renamed for crash safety.
24fn postcard_save_atomic<T: Serialize>(
25    dir: &std::path::Path,
26    filename: &str,
27    value: &T,
28    label: &str,
29) -> Result<(), Error> {
30    let data = postcard::to_allocvec(value).map_err(|e| {
31        Error::Io(std::io::Error::new(
32            std::io::ErrorKind::InvalidData,
33            format!("failed to serialize {label}: {e}"),
34        ))
35    })?;
36    let tmp_path = dir.join(format!("{filename}.tmp"));
37    let final_path = dir.join(filename);
38    std::fs::write(&tmp_path, &data)?;
39    std::fs::rename(&tmp_path, &final_path)?;
40    Ok(())
41}
42
43/// RF-2: Loads and deserializes a postcard file from `dir/filename`.
44///
45/// Returns `Ok(None)` when the file does not exist.
46fn postcard_load<T: for<'de> Deserialize<'de>>(
47    dir: &std::path::Path,
48    filename: &str,
49    label: &str,
50) -> Result<Option<T>, Error> {
51    let path = dir.join(filename);
52    if !path.exists() {
53        return Ok(None);
54    }
55    let file_len = std::fs::metadata(&path)?.len();
56    if file_len > MAX_PQ_ARTIFACT_BYTES {
57        return Err(Error::IndexCorrupted(format!(
58            "{label} file is {file_len} bytes, exceeds cap {MAX_PQ_ARTIFACT_BYTES}"
59        )));
60    }
61    let data = std::fs::read(&path)?;
62    let value: T = postcard::from_bytes(&data).map_err(|e| {
63        Error::Io(std::io::Error::new(
64            std::io::ErrorKind::InvalidData,
65            format!("failed to deserialize {label}: {e}"),
66        ))
67    })?;
68    Ok(Some(value))
69}
70
71/// Persistence methods for codebook and rotation matrix storage.
72impl ProductQuantizer {
73    /// Save trained codebook to `<dir>/codebook.pq` using postcard.
74    /// Uses atomic write (write to .tmp, then rename).
75    ///
76    /// # Errors
77    ///
78    /// Returns `Error::Io` if serialization or file I/O fails.
79    pub fn save_codebook(&self, dir: &std::path::Path) -> Result<(), Error> {
80        postcard_save_atomic(dir, "codebook.pq", self, "PQ codebook")
81    }
82
83    /// Load codebook from `<dir>/codebook.pq`. Returns `None` if file doesn't exist.
84    ///
85    /// The decoded quantizer is structurally validated ([`Self::validate_loaded`])
86    /// before being returned, so a corrupt or tampered codebook is rejected here
87    /// rather than producing out-of-bounds indexing during search.
88    ///
89    /// # Errors
90    ///
91    /// Returns `Error::Io` if deserialization or file I/O fails, or
92    /// `Error::IndexCorrupted` if the decoded codebook/rotation is inconsistent
93    /// or the file exceeds `MAX_PQ_ARTIFACT_BYTES`.
94    pub fn load_codebook(dir: &std::path::Path) -> Result<Option<Self>, Error> {
95        let Some(quantizer): Option<Self> = postcard_load(dir, "codebook.pq", "PQ codebook")?
96        else {
97            return Ok(None);
98        };
99        quantizer.validate_loaded()?;
100        Ok(Some(quantizer))
101    }
102
103    /// Save OPQ rotation matrix to `<dir>/rotation.opq` using postcard.
104    ///
105    /// # Errors
106    ///
107    /// Returns `Error::Io` if the rotation is `None`, serialization, or file I/O fails.
108    pub fn save_rotation(&self, dir: &std::path::Path) -> Result<(), Error> {
109        let rotation = self.rotation.as_ref().ok_or_else(|| {
110            Error::Io(std::io::Error::new(
111                std::io::ErrorKind::InvalidData,
112                "no rotation matrix to save",
113            ))
114        })?;
115        postcard_save_atomic(dir, "rotation.opq", rotation, "OPQ rotation")
116    }
117
118    /// Load OPQ rotation matrix from `<dir>/rotation.opq`. Returns `None` if file doesn't exist.
119    ///
120    /// # Errors
121    ///
122    /// Returns `Error::Io` if deserialization or file I/O fails.
123    pub fn load_rotation(dir: &std::path::Path) -> Result<Option<Vec<f32>>, Error> {
124        postcard_load(dir, "rotation.opq", "OPQ rotation")
125    }
126}
velesdb_core/quantization/pq_persistence.rs

velesdb_core/quantization/
pq_persistence.rs
