Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677use 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 {}. please restart the indexer to use the new dictionary.", dict_bytes.len(), path.display() );
Ok(()) }}