use crate::error::Error;
use serde::{Deserialize, Serialize};
use super::pq::ProductQuantizer;
fn postcard_save_atomic<T: Serialize>(
dir: &std::path::Path,
filename: &str,
value: &T,
label: &str,
) -> Result<(), Error> {
let data = postcard::to_allocvec(value).map_err(|e| {
Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("failed to serialize {label}: {e}"),
))
})?;
let tmp_path = dir.join(format!("{filename}.tmp"));
let final_path = dir.join(filename);
std::fs::write(&tmp_path, &data)?;
std::fs::rename(&tmp_path, &final_path)?;
Ok(())
}
fn postcard_load<T: for<'de> Deserialize<'de>>(
dir: &std::path::Path,
filename: &str,
label: &str,
) -> Result<Option<T>, Error> {
let path = dir.join(filename);
if !path.exists() {
return Ok(None);
}
let data = std::fs::read(&path)?;
let value: T = postcard::from_bytes(&data).map_err(|e| {
Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("failed to deserialize {label}: {e}"),
))
})?;
Ok(Some(value))
}
impl ProductQuantizer {
pub fn save_codebook(&self, dir: &std::path::Path) -> Result<(), Error> {
postcard_save_atomic(dir, "codebook.pq", self, "PQ codebook")
}
pub fn load_codebook(dir: &std::path::Path) -> Result<Option<Self>, Error> {
postcard_load(dir, "codebook.pq", "PQ codebook")
}
pub fn save_rotation(&self, dir: &std::path::Path) -> Result<(), Error> {
let rotation = self.rotation.as_ref().ok_or_else(|| {
Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"no rotation matrix to save",
))
})?;
postcard_save_atomic(dir, "rotation.opq", rotation, "OPQ rotation")
}
pub fn load_rotation(dir: &std::path::Path) -> Result<Option<Vec<f32>>, Error> {
postcard_load(dir, "rotation.opq", "OPQ rotation")
}
}