Improve index BTree insertion & search performance (surrealdb#2240)

lib/Cargo.toml

@@ -78,6 +78,7 @@ indxdb = { version = "0.3.0", optional = true }
 js = { version = "0.3.1" , package = "rquickjs", features = ["array-buffer", "bindgen", "classes", "futures", "loader", "macro", "parallel", "properties","rust-alloc"], optional = true }
 jsonwebtoken = "8.3.0"
 lexicmp = "0.1.0"
+lru = "0.10.1"
 md-5 = "0.10.5"
 nanoid = "0.4.0"
 native-tls = { version = "0.2.11", optional = true }
@@ -150,4 +151,8 @@ harness = false
 
 [[bench]]
 name = "processor"
+harness = false
+
+[[bench]]
+name = "index_btree"
 harness = false

lib/benches/index_btree.rs

@@ -0,0 +1,72 @@
+use criterion::async_executor::FuturesExecutor;
+use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
+use rand::prelude::SliceRandom;
+use rand::thread_rng;
+use serde::de::DeserializeOwned;
+use serde::Serialize;
+use std::time::Duration;
+use surrealdb::idx::bkeys::{BKeys, FstKeys, TrieKeys};
+use surrealdb::idx::btree::store::{BTreeNodeStore, BTreeStoreType, KeyProvider};
+use surrealdb::idx::btree::{BTree, Payload, State};
+use surrealdb::kvs::{Datastore, Key};
+
+macro_rules! get_key_value {
+	($idx:expr) => {{
+		(format!("{}", $idx).into(), ($idx * 10) as Payload)
+	}};
+}
+
+fn bench_index_btree(c: &mut Criterion) {
+	let (samples_len, samples) = setup();
+
+	let mut group = c.benchmark_group("index_btree");
+	group.throughput(Throughput::Elements(1));
+	group.sample_size(10);
+	group.measurement_time(Duration::from_secs(30));
+
+	group.bench_function("btree-insertion-fst", |b| {
+		b.to_async(FuturesExecutor)
+			.iter(|| bench::<_, FstKeys>(samples_len, |i| get_key_value!(samples[i])))
+	});
+
+	group.bench_function("btree-insertion-trie", |b| {
+		b.to_async(FuturesExecutor)
+			.iter(|| bench::<_, TrieKeys>(samples_len, |i| get_key_value!(samples[i])))
+	});
+
+	group.finish();
+}
+
+fn setup() -> (usize, Vec<usize>) {
+	let samples_len = if cfg!(debug_assertions) {
+		1000 // debug is much slower!
+	} else {
+		100_000
+	};
+	let mut samples: Vec<usize> = (0..samples_len).collect();
+	let mut rng = thread_rng();
+	samples.shuffle(&mut rng);
+	(samples_len, samples)
+}
+
+async fn bench<F, BK>(samples_size: usize, sample_provider: F)
+where
+	F: Fn(usize) -> (Key, Payload),
+	BK: BKeys + Serialize + DeserializeOwned + Default,
+{
+	let ds = Datastore::new("memory").await.unwrap();
+	let mut tx = ds.transaction(true, false).await.unwrap();
+	let mut t = BTree::<BK>::new(State::new(100));
+	let s = BTreeNodeStore::new(KeyProvider::Debug, BTreeStoreType::Write, 20);
+	let mut s = s.lock().await;
+	for i in 0..samples_size {
+		let (key, payload) = sample_provider(i);
+		// Insert the sample
+		t.insert(&mut tx, &mut s, key.clone(), payload).await.unwrap();
+		// Search for it
+		black_box(t.search(&mut tx, &mut s, &key).await.unwrap());
+	}
+}
+
+criterion_group!(benches, bench_index_btree);
+criterion_main!(benches);

lib/src/doc/index.rs

@@ -3,6 +3,7 @@ use crate::dbs::Statement;
 use crate::dbs::{Options, Transaction};
 use crate::doc::{CursorDoc, Document};
 use crate::err::Error;
+use crate::idx::btree::store::BTreeStoreType;
 use crate::idx::ft::FtIndex;
 use crate::idx::IndexKeyBase;
 use crate::sql::array::Array;
@@ -192,12 +193,12 @@ impl<'a> IndexOperation<'a> {
 	) -> Result<(), Error> {
 		let ikb = IndexKeyBase::new(self.opt, self.ix);
 		let az = run.get_az(self.opt.ns(), self.opt.db(), az.as_str()).await?;
-		let mut ft = FtIndex::new(run, az, ikb, order, scoring, hl).await?;
+		let mut ft = FtIndex::new(run, az, ikb, order, scoring, hl, BTreeStoreType::Write).await?;
 		if let Some(n) = &self.n {
-			// TODO: Apply the analyzer
-			ft.index_document(run, self.rid, n).await
+			ft.index_document(run, self.rid, n).await?;
 		} else {
-			ft.remove_document(run, self.rid).await
+			ft.remove_document(run, self.rid).await?;
 		}
+		ft.finish(run).await
 	}
 }