chore: add save and get collection bench (#92)

Author: edwinkys

bench/main.rs

@@ -70,4 +70,40 @@ criterion_group!(
     bench_insert_to_collection
 );
 
-criterion_main!(collection);
+fn bench_save_collection_to_database(criterion: &mut Criterion) {
+    let id = "save collection to database";
+
+    // Setup the database and collection.
+    let collection = build_test_collection(DIMENSION, COLLECTION_SIZE);
+    let mut db = create_test_database(DIMENSION, COLLECTION_SIZE);
+
+    // Benchmark the save speed.
+    criterion.bench_function(id, |bencher| {
+        bencher.iter(|| {
+            black_box(db.save_collection("bench", &collection).unwrap());
+        })
+    });
+}
+
+fn bench_get_collection_from_database(criterion: &mut Criterion) {
+    let id = "get collection from database";
+    let db = create_test_database(DIMENSION, COLLECTION_SIZE);
+
+    // Benchmark the get speed.
+    // This is the operation that loads the collection into memory.
+    let routine = || {
+        black_box(db.get_collection("bench").unwrap());
+    };
+
+    criterion.bench_function(id, |b| b.iter(routine));
+}
+
+criterion_group! {
+    name = database;
+    config = Criterion::default().sample_size(10);
+    targets =
+        bench_save_collection_to_database,
+        bench_get_collection_from_database
+}
+
+criterion_main!(collection, database);

bench/utils.rs

@@ -8,3 +8,13 @@ pub fn build_test_collection(dimension: usize, len: usize) -> Collection {
     let config = Config::default();
     Collection::build(&config, &records).unwrap()
 }
+
+/// Creates a pre-populated database with a collection for testing.
+/// * `dimension`: Dimensionality of the vector embeddings
+/// * `size`: Number of records in the collection
+pub fn create_test_database(dimension: usize, size: usize) -> Database {
+    let collection = build_test_collection(dimension, size);
+    let mut db = Database::new("data/bench").unwrap();
+    db.save_collection("bench", &collection).unwrap();
+    db
+}