corenn 0.3.0

Billion-scale vector database
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

use clap::Args;
use libcorenn::store::schema::NODE;
use libcorenn::CoreNN;
use std::path::PathBuf;
use tokio::fs::create_dir_all;
use tokio::fs::File;
use tokio::io::AsyncWriteExt;
use tokio::io::BufWriter;

#[derive(Args)]
pub struct ExportVectorsArgs {
  /// Path to a CoreNN DB.
  #[arg()]
  path: PathBuf,

  /// Output directory to write CoreNN index to.
  #[arg(long)]
  out: PathBuf,
}

impl ExportVectorsArgs {
  pub async fn exec(self) {
    let corenn = CoreNN::open(self.path);

    let dim = corenn.cfg().dim;

    create_dir_all(&self.out).await.unwrap();

    let out_ids = File::create(self.out.join("ids.bin")).await.unwrap();
    let mut out_ids = BufWriter::new(out_ids);
    let out_vecs = File::create(self.out.join("vecs.bin")).await.unwrap();
    let mut out_vecs = BufWriter::new(out_vecs);

    for (id, node) in NODE.iter(&corenn.internal_db()) {
      out_ids.write_u64_le(id.try_into().unwrap()).await.unwrap();
      assert_eq!(node.vector.dim(), dim);
      out_vecs
        .write_all(node.vector.as_raw_bytes())
        .await
        .unwrap();
    }

    out_ids.flush().await.unwrap();
    out_vecs.flush().await.unwrap();

    tracing::info!("all done!");
  }
}