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

velesdb_core/quantization/
pq_persistence.rs
