From 0f2ad59e59e10081b15f6fc2ac2e6a7852cc7ecd Mon Sep 17 00:00:00 2001 From: Jiahao XU Date: Sat, 24 Jul 2021 17:38:07 +1000 Subject: [PATCH 1/2] Add new rust binary vtable.rs Signed-off-by: Jiahao XU --- Cargo.lock | 1 + Cargo.toml | 2 +- bench.sh | 6 ++ src/bin/common.rs | 11 +++ src/bin/vtable.rs | 172 ++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 191 insertions(+), 1 deletion(-) create mode 100644 src/bin/vtable.rs diff --git a/Cargo.lock b/Cargo.lock index 13942ab..737e32c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -832,6 +832,7 @@ dependencies = [ "fallible-iterator", "fallible-streaming-iterator", "hashlink 0.7.0", + "lazy_static", "libsqlite3-sys", "memchr", "smallvec", diff --git a/Cargo.toml b/Cargo.toml index b1aa9a8..b937433 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,4 +10,4 @@ sqlx = { version = "0.5.2", features = ["runtime-tokio-native-tls", "sqlite"]} tokio = {version = "1.5.0", features = ["full"]} rand = "0.8.3" num_cpus = "1.0" -rusqlite = "0.25.3" \ No newline at end of file +rusqlite = { version = "0.25.3", features = ["vtab"] } diff --git a/bench.sh b/bench.sh index fa12f10..aa70f49 100755 --- a/bench.sh +++ b/bench.sh @@ -121,3 +121,9 @@ rm -rf threaded_batched.db threaded_batched.db-shm threaded_batched.db-wal cargo build --release --quiet --bin threaded_batched echo "$(date)" "[RUST] threaded_batched.rs (100_000_000) inserts" time ./target/release/threaded_batched + +# benching where the random generator is exported as a virtual table in sqlite +rm -rf vtable.db +cargo build --release --quiet --bin vtable +echo "$(date)" "[RUST] vtable.rs (100_000_000) inserts" +time ./target/release/vtable diff --git a/src/bin/common.rs b/src/bin/common.rs index 8cd4e92..923eed3 100644 --- a/src/bin/common.rs +++ b/src/bin/common.rs @@ -21,3 +21,14 @@ pub fn get_random_area_code() -> String { let mut rng = rand::thread_rng(); format!("{:06}", rng.gen_range(0..999999)) } + +pub fn get_random_area_code_u8() -> [u8; 6] { + let mut rng = rand::thread_rng(); + + let mut ret: [u8; 6] = Default::default(); + for each in &mut ret { + *each = b'0' + rng.gen_range(0..9); + } + + ret +} diff --git a/src/bin/vtable.rs b/src/bin/vtable.rs new file mode 100644 index 0000000..8933450 --- /dev/null +++ b/src/bin/vtable.rs @@ -0,0 +1,172 @@ +use std::os::raw::c_int; +use std::marker::PhantomData; +use std::mem::MaybeUninit; +use std::str::from_utf8_unchecked; + +use rusqlite::{Connection, Result, Error, params}; +use rusqlite::vtab::{ + Context, Values, + VTab, VTabConnection, IndexInfo, sqlite3_vtab, + VTabCursor, sqlite3_vtab_cursor, + eponymous_only_module, +}; + +mod common; + +#[repr(C)] +struct RandVTable { + /// Base class. Must be first + base: sqlite3_vtab, + /* Virtual table implementations will typically add additional fields */ +} +impl RandVTable { + fn register(conn: &Connection) -> Result<()> { + conn.create_module( + "rand_vtab", + eponymous_only_module::<'_, Self>(), + None + ) + } +} +unsafe impl<'vtab> VTab<'vtab> for RandVTable { + type Aux = (); + type Cursor = RandVTableCursor<'vtab>; + + fn connect( + _db: &mut VTabConnection, + _aux: Option<&Self::Aux>, + _args: &[&[u8]] + ) -> Result<(String, Self)> { + Ok(( + "CREATE TABLE user ( + area CHAR(6), + age INTEGER not null, + active INTEGER not null + )".to_owned(), + Self { base: Default::default() } + )) + } + + fn best_index(&self, info: &mut IndexInfo) -> Result<()> { + for i in 0..info.constraints().count() { + let mut index_constraint_usage = info.constraint_usage(i); + // no need of constrain in `::filter`. + index_constraint_usage.set_argv_index(0); + // VTabCursor does not test for constrain, so sqlite3 must not + // omit the test. + index_constraint_usage.set_omit(false); + } + // RandVTable does not return ordered rows + info.set_order_by_consumed(false); + + // idx_num is unused + info.set_idx_num(0); + + // RandVTable has infinite tables + let estimated_rows = i64::MAX; + //info.set_estimated_rows(estimated_rows); + + // estimated_cost is linear + info.set_estimated_cost(estimated_rows as f64); + + Ok(()) + } + + fn open(&'vtab self) -> Result { + Ok(Self::Cursor::new()) + } +} + +#[repr(C)] +struct RandVTableCursor<'vtab> { + /// Base class. Must be first + base: sqlite3_vtab_cursor, + /* Virtual table implementations will typically add additional fields */ + phantom: PhantomData<&'vtab RandVTable>, + + rowid: u64, + + area_code: [u8; 6], + age: i8, + active: i8, +} +impl<'vtab> RandVTableCursor<'vtab> { + fn new() -> Self { + Self { + base: unsafe { MaybeUninit::zeroed().assume_init() }, + phantom: PhantomData, + rowid: 0, + area_code: common::get_random_area_code_u8(), + age: common::get_random_age(), + active: common::get_random_active(), + } + } +} +unsafe impl<'vtab> VTabCursor for RandVTableCursor<'vtab> { + /// RandVTableCursor doesn't need any filter capacity + fn filter( + &mut self, + _idx_num: c_int, + _idx_str: Option<&str>, + _args: &Values<'_> + ) -> Result<()> { + Ok(()) + } + + fn next(&mut self) -> Result<()> { + self.rowid += 1; + + self.area_code = common::get_random_area_code_u8(); + self.age = common::get_random_age(); + self.active = common::get_random_active(); + + Ok(()) + } + + /// RandVTableCursor is endless + fn eof(&self) -> bool { + false + } + + fn column(&self, ctx: &mut Context, i: c_int) -> Result<()> { + match i { + 0 => ctx.set_result(unsafe { &from_utf8_unchecked(&self.area_code) })?, + 1 => ctx.set_result(&self.age)?, + 2 => ctx.set_result(&self.active)?, + _ => return Err(Error::InvalidColumnIndex(i as usize)), + }; + Ok(()) + } + + fn rowid(&self) -> Result { + Ok(self.rowid as i64) + } +} + +fn faker(mut conn: Connection, count: i64) { + let tx = conn.transaction().unwrap(); + tx.execute( + "INSERT INTO user(area, age, active) + SELECT * FROM rand_vtab LIMIT ?", + params![count], + ).unwrap(); + tx.commit().unwrap(); +} + + +fn main() { + let conn = Connection::open("vtable.db").unwrap(); + conn.execute( + "CREATE TABLE IF NOT EXISTS user ( + id INTEGER not null primary key, + area CHAR(6), + age INTEGER not null, + active INTEGER not null + )", + [], + ).unwrap(); + + RandVTable::register(&conn).unwrap(); + + faker(conn, 100_000_000) +} From 09d819852a3eba740e054b14f0ebab71482c38c2 Mon Sep 17 00:00:00 2001 From: Jiahao XU Date: Sat, 24 Jul 2021 21:37:46 +1000 Subject: [PATCH 2/2] Fix vtable::RandVTableCursor: Support for Null area_code Signed-off-by: Jiahao XU --- src/bin/common.rs | 8 ++++++++ src/bin/vtable.rs | 15 +++++++++++---- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/bin/common.rs b/src/bin/common.rs index 923eed3..83baed1 100644 --- a/src/bin/common.rs +++ b/src/bin/common.rs @@ -32,3 +32,11 @@ pub fn get_random_area_code_u8() -> [u8; 6] { ret } + +pub fn get_random_optional_area_code_u8() -> Option<[u8; 6]> { + if get_random_bool() { + Some(get_random_area_code_u8()) + } else { + None + } +} diff --git a/src/bin/vtable.rs b/src/bin/vtable.rs index 8933450..5036b6a 100644 --- a/src/bin/vtable.rs +++ b/src/bin/vtable.rs @@ -4,6 +4,7 @@ use std::mem::MaybeUninit; use std::str::from_utf8_unchecked; use rusqlite::{Connection, Result, Error, params}; +use rusqlite::types::Null; use rusqlite::vtab::{ Context, Values, VTab, VTabConnection, IndexInfo, sqlite3_vtab, @@ -86,7 +87,7 @@ struct RandVTableCursor<'vtab> { rowid: u64, - area_code: [u8; 6], + area_code: Option<[u8; 6]>, age: i8, active: i8, } @@ -96,7 +97,7 @@ impl<'vtab> RandVTableCursor<'vtab> { base: unsafe { MaybeUninit::zeroed().assume_init() }, phantom: PhantomData, rowid: 0, - area_code: common::get_random_area_code_u8(), + area_code: common::get_random_optional_area_code_u8(), age: common::get_random_age(), active: common::get_random_active(), } @@ -116,7 +117,7 @@ unsafe impl<'vtab> VTabCursor for RandVTableCursor<'vtab> { fn next(&mut self) -> Result<()> { self.rowid += 1; - self.area_code = common::get_random_area_code_u8(); + self.area_code = common::get_random_optional_area_code_u8(); self.age = common::get_random_age(); self.active = common::get_random_active(); @@ -130,7 +131,13 @@ unsafe impl<'vtab> VTabCursor for RandVTableCursor<'vtab> { fn column(&self, ctx: &mut Context, i: c_int) -> Result<()> { match i { - 0 => ctx.set_result(unsafe { &from_utf8_unchecked(&self.area_code) })?, + 0 => { + if let Some(area_code) = self.area_code { + ctx.set_result(unsafe { &from_utf8_unchecked(&area_code) })?; + } else { + ctx.set_result(&Null)?; + } + }, 1 => ctx.set_result(&self.age)?, 2 => ctx.set_result(&self.active)?, _ => return Err(Error::InvalidColumnIndex(i as usize)),