Export the random generation code as a virtual table in sqlite3

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"
+rusqlite = { version = "0.25.3", features = ["vtab"] }

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

src/bin/common.rs

@@ -21,3 +21,22 @@ 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
+}
+
+pub fn get_random_optional_area_code_u8() -> Option<[u8; 6]> {
+    if get_random_bool() {
+        Some(get_random_area_code_u8())
+    } else {
+        None
+    }
+}

src/bin/vtable.rs

@@ -0,0 +1,179 @@
+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::types::Null;
+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 `<RandVTableCursor as VTabCursor>::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<Self::Cursor> {
+        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: Option<[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_optional_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_optional_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 => {
+                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)),
+        };
+        Ok(())
+    }
+
+    fn rowid(&self) -> Result<i64> {
+        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)
+}