use miette::{IntoDiagnostic, Result};
use std::cell::RefCell;
use std::collections::HashSet;
use super::{Db, registry};
impl Db {
pub fn train_dict(&self, ks_name: &str) -> Result<()> {
let Some((name, cfg)) = registry::trainable()
.into_iter()
.find(|(name, _)| *name == ks_name)
else {
miette::bail!("unknown keyspace for training: {ks_name}");
};
let ks = self
.keyspace_by_name(name)
.expect("trainable keyspace is registered");
tracing::info!("collecting training samples for keyspace {ks_name}...");
let mut seen_keys = HashSet::new();
let captured_keys = RefCell::new(Vec::new());
let progress = RefCell::new((std::time::Instant::now(), 0));
let samples = ks
.sample_data_blocks(cfg.samples, |first, last| {
let passes = !seen_keys.contains(&(first.to_vec(), last.to_vec()));
if passes {
captured_keys
.borrow_mut()
.push((first.to_vec(), last.to_vec()));
}
let mut p = progress.borrow_mut();
if passes {
p.1 += 1;
}
if p.0.elapsed() >= std::time::Duration::from_secs(5) {
tracing::info!(
"collected {}/{} samples for {}...",
p.1,
cfg.samples,
ks_name
);
p.0 = std::time::Instant::now();
}
passes
})
.into_diagnostic()?
.into_iter()
.zip(captured_keys.into_inner())
.filter_map(|(sample, keys)| seen_keys.insert(keys).then_some(sample.to_vec()))
.collect::<Vec<_>>();
if samples.is_empty() {
miette::bail!("no samples found for keyspace {ks_name}, skipping training");
}
tracing::info!(
"training zstd dictionary for keyspace {ks_name} ({} samples, {} bytes limit)...",
samples.len(),
cfg.dict_size,
);
let dict_bytes =
lsm_tree::train_zstd_dict(&samples, cfg.dict_size as usize).into_diagnostic()?;
let path = self.path.join(format!("dict_{ks_name}.bin"));
std::fs::write(&path, &dict_bytes).into_diagnostic()?;
tracing::info!(
"saved zstd dictionary for keyspace {ks_name} ({} bytes) to {}. only keyspaces created from now on use it, existing ones keep their compression",
dict_bytes.len(),
path.display()
);
Ok(())
}
}