mirror of
https://github.com/tursodatabase/limbo.git
synced 2025-08-04 18:18:03 +00:00
Merge 'Implement create virtual table using vtab modules, more work on virtual tables' from Preston Thorpe
This PR started out as one to improve the API of extensions but I ended up building on top of this quite a bit and it just kept going. Sorry this one is so large but there wasn't really a good stopping point, as it kept leaving stuff in broken states. **VCreate**: Support for `CREATE VIRTUAL TABLE t USING vtab_module` **VUpdate**: Support for `INSERT` and `DELETE` methods on virtual tables. Sqlite uses `xUpdate` function with the `VUpdate` opcode to handle all insert/update/delete functionality in virtual tables.. have to just document that: ``` if args[0] == NULL: INSERT args[1] the values in args[2..] if args[1] == NULL: DELETE args[0] if args[0] != NULL && len(args) > 2: Update values=args[2..] rowid=args[0] ``` I know I asked @jussisaurio on discord about this already, but it just sucked so bad that I added some internal translation so we could expose a [nice API](https://github.com/tursodatabase/limbo/pull/996/files#diff- 3e8f8a660b11786745b48b528222d11671e9f19fa00a032a4eefb5412e8200d1R54) and handle the logic ourselves while keeping with sqlite's opcodes. I'll change it back if I have to, I just thought it was genuinely awful to have to rely on comments to explain all that to extension authors. The included extension is not meant to be a legitimately useful one, it is there for testing purposes. I did something similar in #960 using a test extension, so I figure when they are both merged, I will go back and combine them into one since you can do many kinds at once, and that way it will reduce the amount of crates and therefore compile time. 1. Remaining opcodes. 2. `UPDATE` (when we support the syntax) 3. `xConnect` - expose API for a DB connection to a vtab so it can perform arbitrary queries. Closes #996
This commit is contained in:
Pekka Enberg
commit
7f2525ac27
30 changed files with 1362 additions and 378 deletions
|
|
@ -95,6 +95,9 @@ impl AggFunc for Percentile {
|
|||
/// The state to track during the steps
|
||||
type State = (Vec<f64>, Option<f64>, Option<String>); // Tracks the values, Percentile, and errors
|
||||
|
||||
/// Define your error type, must impl Display
|
||||
type Error = String;
|
||||
|
||||
/// Define the name you wish to call your function by.
|
||||
/// e.g. SELECT percentile(value, 40);
|
||||
const NAME: &str = "percentile";
|
||||
|
|
@ -129,15 +132,15 @@ impl AggFunc for Percentile {
|
|||
}
|
||||
/// A function to finalize the state into a value to be returned as a result
|
||||
/// or an error (if you chose to track an error state as well)
|
||||
fn finalize(state: Self::State) -> Value {
|
||||
fn finalize(state: Self::State) -> Result<Value, Self::Error> {
|
||||
let (mut values, p_value, error) = state;
|
||||
|
||||
if let Some(error) = error {
|
||||
return Value::custom_error(error);
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
if values.is_empty() {
|
||||
return Value::null();
|
||||
return Ok(Value::null());
|
||||
}
|
||||
|
||||
values.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
||||
|
|
@ -145,7 +148,7 @@ impl AggFunc for Percentile {
|
|||
let p = p_value.unwrap();
|
||||
let index = (p * (n - 1.0) / 100.0).floor() as usize;
|
||||
|
||||
Value::from_float(values[index])
|
||||
Ok(Value::from_float(values[index]))
|
||||
}
|
||||
}
|
||||
```
|
||||
|
|
@ -161,21 +164,25 @@ struct CsvVTable;
|
|||
|
||||
impl VTabModule for CsvVTable {
|
||||
type VCursor = CsvCursor;
|
||||
/// Define your error type. Must impl Display and match VCursor::Error
|
||||
type Error = &'static str;
|
||||
/// Declare the name for your virtual table
|
||||
const NAME: &'static str = "csv_data";
|
||||
|
||||
/// Declare the table schema and call `api.declare_virtual_table` with the schema sql.
|
||||
fn connect(api: &ExtensionApi) -> ResultCode {
|
||||
let sql = "CREATE TABLE csv_data(
|
||||
/// Declare the type of vtable (TableValuedFunction or VirtualTable)
|
||||
const VTAB_KIND: VTabKind = VTabKind::VirtualTable;
|
||||
|
||||
/// Function to initialize the schema of your vtable
|
||||
fn create_schema(_args: &[Value]) -> &'static str {
|
||||
"CREATE TABLE csv_data(
|
||||
name TEXT,
|
||||
age TEXT,
|
||||
city TEXT
|
||||
)";
|
||||
api.declare_virtual_table(Self::NAME, sql)
|
||||
)"
|
||||
}
|
||||
|
||||
/// Open to return a new cursor: In this simple example, the CSV file is read completely into memory on connect.
|
||||
fn open() -> Self::VCursor {
|
||||
fn open(&self) -> Result<Self::VCursor, Self::Error> {
|
||||
// Read CSV file contents from "data.csv"
|
||||
let csv_content = fs::read_to_string("data.csv").unwrap_or_default();
|
||||
// For simplicity, we'll ignore the header row.
|
||||
|
|
@ -188,16 +195,16 @@ impl VTabModule for CsvVTable {
|
|||
.collect()
|
||||
})
|
||||
.collect();
|
||||
CsvCursor { rows, index: 0 }
|
||||
Ok(CsvCursor { rows, index: 0 })
|
||||
}
|
||||
|
||||
/// Filter through result columns. (not used in this simple example)
|
||||
fn filter(_cursor: &mut Self::VCursor, _arg_count: i32, _args: &[Value]) -> ResultCode {
|
||||
fn filter(_cursor: &mut Self::VCursor, _args: &[Value]) -> ResultCode {
|
||||
ResultCode::OK
|
||||
}
|
||||
|
||||
/// Return the value for the column at the given index in the current row.
|
||||
fn column(cursor: &Self::VCursor, idx: u32) -> Value {
|
||||
fn column(cursor: &Self::VCursor, idx: u32) -> Result<Value, Self::Error> {
|
||||
cursor.column(idx)
|
||||
}
|
||||
|
||||
|
|
@ -215,6 +222,22 @@ impl VTabModule for CsvVTable {
|
|||
fn eof(cursor: &Self::VCursor) -> bool {
|
||||
cursor.index >= cursor.rows.len()
|
||||
}
|
||||
|
||||
/// *Optional* methods for non-readonly tables
|
||||
|
||||
/// Update the value at rowid
|
||||
fn update(&mut self, _rowid: i64, _args: &[Value]) -> Result<(), Self::Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Insert the value(s)
|
||||
fn insert(&mut self, _args: &[Value]) -> Result<i64, Self::Error> {
|
||||
Ok(0)
|
||||
}
|
||||
/// Delete the value at rowid
|
||||
fn delete(&mut self, _rowid: i64) -> Result<(), Self::Error> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// The cursor for iterating over CSV rows.
|
||||
|
|
@ -226,6 +249,8 @@ struct CsvCursor {
|
|||
|
||||
/// Implement the VTabCursor trait for your cursor type
|
||||
impl VTabCursor for CsvCursor {
|
||||
type Error = &'static str;
|
||||
|
||||
fn next(&mut self) -> ResultCode {
|
||||
CsvCursor::next(self)
|
||||
}
|
||||
|
|
@ -234,12 +259,12 @@ impl VTabCursor for CsvCursor {
|
|||
self.index >= self.rows.len()
|
||||
}
|
||||
|
||||
fn column(&self, idx: u32) -> Value {
|
||||
fn column(&self, idx: u32) -> Result<Value, Self::Error> {
|
||||
let row = &self.rows[self.index];
|
||||
if (idx as usize) < row.len() {
|
||||
Value::from_text(&row[idx as usize])
|
||||
Ok(Value::from_text(&row[idx as usize]))
|
||||
} else {
|
||||
Value::null()
|
||||
Ok(Value::null())
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,63 +1,48 @@
|
|||
mod types;
|
||||
pub use limbo_macros::{register_extension, scalar, AggregateDerive, VTabModuleDerive};
|
||||
use std::os::raw::{c_char, c_void};
|
||||
use std::{
|
||||
fmt::Display,
|
||||
os::raw::{c_char, c_void},
|
||||
};
|
||||
pub use types::{ResultCode, Value, ValueType};
|
||||
|
||||
pub type ExtResult<T> = std::result::Result<T, ResultCode>;
|
||||
|
||||
#[repr(C)]
|
||||
pub struct ExtensionApi {
|
||||
pub ctx: *mut c_void,
|
||||
|
||||
pub register_scalar_function: unsafe extern "C" fn(
|
||||
ctx: *mut c_void,
|
||||
name: *const c_char,
|
||||
func: ScalarFunction,
|
||||
) -> ResultCode,
|
||||
|
||||
pub register_aggregate_function: unsafe extern "C" fn(
|
||||
ctx: *mut c_void,
|
||||
name: *const c_char,
|
||||
args: i32,
|
||||
init_func: InitAggFunction,
|
||||
step_func: StepFunction,
|
||||
finalize_func: FinalizeFunction,
|
||||
) -> ResultCode,
|
||||
|
||||
pub register_module: unsafe extern "C" fn(
|
||||
ctx: *mut c_void,
|
||||
name: *const c_char,
|
||||
module: VTabModuleImpl,
|
||||
) -> ResultCode,
|
||||
|
||||
pub declare_vtab: unsafe extern "C" fn(
|
||||
ctx: *mut c_void,
|
||||
name: *const c_char,
|
||||
sql: *const c_char,
|
||||
) -> ResultCode,
|
||||
}
|
||||
|
||||
impl ExtensionApi {
|
||||
pub fn declare_virtual_table(&self, name: &str, sql: &str) -> ResultCode {
|
||||
let Ok(name) = std::ffi::CString::new(name) else {
|
||||
return ResultCode::Error;
|
||||
};
|
||||
let Ok(sql) = std::ffi::CString::new(sql) else {
|
||||
return ResultCode::Error;
|
||||
};
|
||||
unsafe { (self.declare_vtab)(self.ctx, name.as_ptr(), sql.as_ptr()) }
|
||||
}
|
||||
pub register_scalar_function: RegisterScalarFn,
|
||||
pub register_aggregate_function: RegisterAggFn,
|
||||
pub register_module: RegisterModuleFn,
|
||||
}
|
||||
|
||||
pub type ExtensionEntryPoint = unsafe extern "C" fn(api: *const ExtensionApi) -> ResultCode;
|
||||
|
||||
pub type ScalarFunction = unsafe extern "C" fn(argc: i32, *const Value) -> Value;
|
||||
|
||||
pub type RegisterScalarFn =
|
||||
unsafe extern "C" fn(ctx: *mut c_void, name: *const c_char, func: ScalarFunction) -> ResultCode;
|
||||
|
||||
pub type RegisterAggFn = unsafe extern "C" fn(
|
||||
ctx: *mut c_void,
|
||||
name: *const c_char,
|
||||
args: i32,
|
||||
init: InitAggFunction,
|
||||
step: StepFunction,
|
||||
finalize: FinalizeFunction,
|
||||
) -> ResultCode;
|
||||
|
||||
pub type RegisterModuleFn = unsafe extern "C" fn(
|
||||
ctx: *mut c_void,
|
||||
name: *const c_char,
|
||||
module: VTabModuleImpl,
|
||||
kind: VTabKind,
|
||||
) -> ResultCode;
|
||||
|
||||
pub type InitAggFunction = unsafe extern "C" fn() -> *mut AggCtx;
|
||||
pub type StepFunction = unsafe extern "C" fn(ctx: *mut AggCtx, argc: i32, argv: *const Value);
|
||||
pub type FinalizeFunction = unsafe extern "C" fn(ctx: *mut AggCtx) -> Value;
|
||||
|
||||
pub trait Scalar {
|
||||
fn call(&self, args: &[Value]) -> Value;
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
pub struct AggCtx {
|
||||
pub state: *mut c_void,
|
||||
|
|
@ -65,59 +50,99 @@ pub struct AggCtx {
|
|||
|
||||
pub trait AggFunc {
|
||||
type State: Default;
|
||||
type Error: Display;
|
||||
const NAME: &'static str;
|
||||
const ARGS: i32;
|
||||
|
||||
fn step(state: &mut Self::State, args: &[Value]);
|
||||
fn finalize(state: Self::State) -> Value;
|
||||
fn finalize(state: Self::State) -> Result<Value, Self::Error>;
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct VTabModuleImpl {
|
||||
pub ctx: *const c_void,
|
||||
pub name: *const c_char,
|
||||
pub connect: VtabFnConnect,
|
||||
pub create_schema: VtabFnCreateSchema,
|
||||
pub open: VtabFnOpen,
|
||||
pub filter: VtabFnFilter,
|
||||
pub column: VtabFnColumn,
|
||||
pub next: VtabFnNext,
|
||||
pub eof: VtabFnEof,
|
||||
pub update: VtabFnUpdate,
|
||||
pub rowid: VtabRowIDFn,
|
||||
}
|
||||
|
||||
pub type VtabFnConnect = unsafe extern "C" fn(api: *const c_void) -> ResultCode;
|
||||
impl VTabModuleImpl {
|
||||
pub fn init_schema(&self, args: Vec<Value>) -> ExtResult<String> {
|
||||
let schema = unsafe { (self.create_schema)(args.as_ptr(), args.len() as i32) };
|
||||
if schema.is_null() {
|
||||
return Err(ResultCode::InvalidArgs);
|
||||
}
|
||||
for arg in args {
|
||||
unsafe { arg.free() };
|
||||
}
|
||||
let schema = unsafe { std::ffi::CString::from_raw(schema) };
|
||||
Ok(schema.to_string_lossy().to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub type VtabFnOpen = unsafe extern "C" fn() -> *mut c_void;
|
||||
pub type VtabFnCreateSchema = unsafe extern "C" fn(args: *const Value, argc: i32) -> *mut c_char;
|
||||
|
||||
pub type VtabFnOpen = unsafe extern "C" fn(*const c_void) -> *const c_void;
|
||||
|
||||
pub type VtabFnFilter =
|
||||
unsafe extern "C" fn(cursor: *mut c_void, argc: i32, argv: *const Value) -> ResultCode;
|
||||
unsafe extern "C" fn(cursor: *const c_void, argc: i32, argv: *const Value) -> ResultCode;
|
||||
|
||||
pub type VtabFnColumn = unsafe extern "C" fn(cursor: *mut c_void, idx: u32) -> Value;
|
||||
pub type VtabFnColumn = unsafe extern "C" fn(cursor: *const c_void, idx: u32) -> Value;
|
||||
|
||||
pub type VtabFnNext = unsafe extern "C" fn(cursor: *mut c_void) -> ResultCode;
|
||||
pub type VtabFnNext = unsafe extern "C" fn(cursor: *const c_void) -> ResultCode;
|
||||
|
||||
pub type VtabFnEof = unsafe extern "C" fn(cursor: *mut c_void) -> bool;
|
||||
pub type VtabFnEof = unsafe extern "C" fn(cursor: *const c_void) -> bool;
|
||||
|
||||
pub type VtabRowIDFn = unsafe extern "C" fn(cursor: *const c_void) -> i64;
|
||||
|
||||
pub type VtabFnUpdate = unsafe extern "C" fn(
|
||||
vtab: *const c_void,
|
||||
argc: i32,
|
||||
argv: *const Value,
|
||||
p_out_rowid: *mut i64,
|
||||
) -> ResultCode;
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub enum VTabKind {
|
||||
VirtualTable,
|
||||
TableValuedFunction,
|
||||
}
|
||||
|
||||
pub trait VTabModule: 'static {
|
||||
type VCursor: VTabCursor;
|
||||
type VCursor: VTabCursor<Error = Self::Error>;
|
||||
const VTAB_KIND: VTabKind;
|
||||
const NAME: &'static str;
|
||||
type Error: std::fmt::Display;
|
||||
|
||||
fn connect(api: &ExtensionApi) -> ResultCode;
|
||||
fn open() -> Self::VCursor;
|
||||
fn filter(cursor: &mut Self::VCursor, arg_count: i32, args: &[Value]) -> ResultCode;
|
||||
fn column(cursor: &Self::VCursor, idx: u32) -> Value;
|
||||
fn create_schema(args: &[Value]) -> String;
|
||||
fn open(&self) -> Result<Self::VCursor, Self::Error>;
|
||||
fn filter(cursor: &mut Self::VCursor, args: &[Value]) -> ResultCode;
|
||||
fn column(cursor: &Self::VCursor, idx: u32) -> Result<Value, Self::Error>;
|
||||
fn next(cursor: &mut Self::VCursor) -> ResultCode;
|
||||
fn eof(cursor: &Self::VCursor) -> bool;
|
||||
fn update(&mut self, _rowid: i64, _args: &[Value]) -> Result<(), Self::Error> {
|
||||
Ok(())
|
||||
}
|
||||
fn insert(&mut self, _args: &[Value]) -> Result<i64, Self::Error> {
|
||||
Ok(0)
|
||||
}
|
||||
fn delete(&mut self, _rowid: i64) -> Result<(), Self::Error> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub trait VTabCursor: Sized {
|
||||
type Error: std::fmt::Display;
|
||||
fn rowid(&self) -> i64;
|
||||
fn column(&self, idx: u32) -> Value;
|
||||
fn column(&self, idx: u32) -> Result<Value, Self::Error>;
|
||||
fn eof(&self) -> bool;
|
||||
fn next(&mut self) -> ResultCode;
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
pub struct VTabImpl {
|
||||
pub module: VTabModuleImpl,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -21,6 +21,8 @@ pub enum ResultCode {
|
|||
Unavailable = 13,
|
||||
CustomError = 14,
|
||||
EOF = 15,
|
||||
ReadOnly = 16,
|
||||
RowID = 17,
|
||||
}
|
||||
|
||||
impl ResultCode {
|
||||
|
|
@ -52,6 +54,8 @@ impl Display for ResultCode {
|
|||
ResultCode::Unavailable => write!(f, "Unavailable"),
|
||||
ResultCode::CustomError => write!(f, "Error "),
|
||||
ResultCode::EOF => write!(f, "EOF"),
|
||||
ResultCode::ReadOnly => write!(f, "Read Only"),
|
||||
ResultCode::RowID => write!(f, "RowID"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -403,6 +407,7 @@ impl Value {
|
|||
}
|
||||
}
|
||||
|
||||
/// Extension authors should __not__ use this function.
|
||||
/// # Safety
|
||||
/// consumes the value while freeing the underlying memory with null check.
|
||||
/// however this does assume that the type was properly constructed with
|
||||
|
|
|
|||
20
extensions/kvstore/Cargo.toml
Normal file
20
extensions/kvstore/Cargo.toml
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
[package]
|
||||
name = "limbo_kv"
|
||||
version.workspace = true
|
||||
authors.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[lib]
|
||||
crate-type = ["cdylib", "lib"]
|
||||
|
||||
[features]
|
||||
static= [ "limbo_ext/static" ]
|
||||
|
||||
[dependencies]
|
||||
lazy_static = "1.5.0"
|
||||
limbo_ext = { workspace = true, features = ["static"] }
|
||||
|
||||
[target.'cfg(not(target_family = "wasm"))'.dependencies]
|
||||
mimalloc = { version = "*", default-features = false }
|
||||
147
extensions/kvstore/src/lib.rs
Normal file
147
extensions/kvstore/src/lib.rs
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
use lazy_static::lazy_static;
|
||||
use limbo_ext::{
|
||||
register_extension, ResultCode, VTabCursor, VTabKind, VTabModule, VTabModuleDerive, Value,
|
||||
};
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Mutex;
|
||||
|
||||
lazy_static! {
|
||||
static ref GLOBAL_STORE: Mutex<BTreeMap<i64, (String, String)>> = Mutex::new(BTreeMap::new());
|
||||
}
|
||||
|
||||
register_extension! {
|
||||
vtabs: { KVStoreVTab },
|
||||
}
|
||||
|
||||
#[derive(VTabModuleDerive, Default)]
|
||||
pub struct KVStoreVTab;
|
||||
|
||||
/// the cursor holds a snapshot of (rowid, key, value) in memory.
|
||||
pub struct KVStoreCursor {
|
||||
rows: Vec<(i64, String, String)>,
|
||||
index: Option<usize>,
|
||||
}
|
||||
|
||||
impl VTabModule for KVStoreVTab {
|
||||
type VCursor = KVStoreCursor;
|
||||
const VTAB_KIND: VTabKind = VTabKind::VirtualTable;
|
||||
const NAME: &'static str = "kv_store";
|
||||
type Error = String;
|
||||
|
||||
fn create_schema(_args: &[Value]) -> String {
|
||||
"CREATE TABLE x (key TEXT PRIMARY KEY, value TEXT);".to_string()
|
||||
}
|
||||
|
||||
fn open(&self) -> Result<Self::VCursor, Self::Error> {
|
||||
Ok(KVStoreCursor {
|
||||
rows: Vec::new(),
|
||||
index: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn filter(cursor: &mut Self::VCursor, _args: &[Value]) -> ResultCode {
|
||||
let store = GLOBAL_STORE.lock().unwrap();
|
||||
cursor.rows = store
|
||||
.iter()
|
||||
.map(|(&rowid, (k, v))| (rowid, k.clone(), v.clone()))
|
||||
.collect();
|
||||
cursor.rows.sort_by_key(|(rowid, _, _)| *rowid);
|
||||
|
||||
if cursor.rows.is_empty() {
|
||||
cursor.index = None;
|
||||
return ResultCode::EOF;
|
||||
} else {
|
||||
cursor.index = Some(0);
|
||||
}
|
||||
ResultCode::OK
|
||||
}
|
||||
|
||||
fn insert(&mut self, values: &[Value]) -> Result<i64, Self::Error> {
|
||||
let key = values
|
||||
.first()
|
||||
.and_then(|v| v.to_text())
|
||||
.ok_or("Missing key")?
|
||||
.to_string();
|
||||
let val = values
|
||||
.get(1)
|
||||
.and_then(|v| v.to_text())
|
||||
.ok_or("Missing value")?
|
||||
.to_string();
|
||||
let rowid = hash_key(&key);
|
||||
{
|
||||
let mut store = GLOBAL_STORE.lock().unwrap();
|
||||
store.insert(rowid, (key, val));
|
||||
}
|
||||
Ok(rowid)
|
||||
}
|
||||
|
||||
fn delete(&mut self, rowid: i64) -> Result<(), Self::Error> {
|
||||
let mut store = GLOBAL_STORE.lock().unwrap();
|
||||
store.remove(&rowid);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn update(&mut self, rowid: i64, values: &[Value]) -> Result<(), Self::Error> {
|
||||
{
|
||||
let mut store = GLOBAL_STORE.lock().unwrap();
|
||||
store.remove(&rowid);
|
||||
}
|
||||
let _ = self.insert(values)?;
|
||||
Ok(())
|
||||
}
|
||||
fn eof(cursor: &Self::VCursor) -> bool {
|
||||
cursor.index.is_some_and(|s| s >= cursor.rows.len()) || cursor.index.is_none()
|
||||
}
|
||||
|
||||
fn next(cursor: &mut Self::VCursor) -> ResultCode {
|
||||
cursor.index = Some(cursor.index.unwrap_or(0) + 1);
|
||||
if cursor.index.is_some_and(|c| c >= cursor.rows.len()) {
|
||||
return ResultCode::EOF;
|
||||
}
|
||||
ResultCode::OK
|
||||
}
|
||||
|
||||
fn column(cursor: &Self::VCursor, idx: u32) -> Result<Value, Self::Error> {
|
||||
if cursor.index.is_some_and(|c| c >= cursor.rows.len()) {
|
||||
return Err("cursor out of range".into());
|
||||
}
|
||||
let (_, ref key, ref val) = cursor.rows[cursor.index.unwrap_or(0)];
|
||||
match idx {
|
||||
0 => Ok(Value::from_text(key.clone())), // key
|
||||
1 => Ok(Value::from_text(val.clone())), // value
|
||||
_ => Err("Invalid column".into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn hash_key(key: &str) -> i64 {
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
key.hash(&mut hasher);
|
||||
hasher.finish() as i64
|
||||
}
|
||||
|
||||
impl VTabCursor for KVStoreCursor {
|
||||
type Error = String;
|
||||
|
||||
fn rowid(&self) -> i64 {
|
||||
if self.index.is_some_and(|c| c < self.rows.len()) {
|
||||
self.rows[self.index.unwrap_or(0)].0
|
||||
} else {
|
||||
println!("rowid: -1");
|
||||
-1
|
||||
}
|
||||
}
|
||||
|
||||
fn column(&self, idx: u32) -> Result<Value, Self::Error> {
|
||||
<KVStoreVTab as VTabModule>::column(self, idx)
|
||||
}
|
||||
|
||||
fn eof(&self) -> bool {
|
||||
<KVStoreVTab as VTabModule>::eof(self)
|
||||
}
|
||||
|
||||
fn next(&mut self) -> ResultCode {
|
||||
<KVStoreVTab as VTabModule>::next(self)
|
||||
}
|
||||
}
|
||||
|
|
@ -9,6 +9,7 @@ struct Median;
|
|||
|
||||
impl AggFunc for Median {
|
||||
type State = Vec<f64>;
|
||||
type Error = &'static str;
|
||||
const NAME: &'static str = "median";
|
||||
const ARGS: i32 = 1;
|
||||
|
||||
|
|
@ -18,9 +19,9 @@ impl AggFunc for Median {
|
|||
}
|
||||
}
|
||||
|
||||
fn finalize(state: Self::State) -> Value {
|
||||
fn finalize(state: Self::State) -> Result<Value, Self::Error> {
|
||||
if state.is_empty() {
|
||||
return Value::null();
|
||||
return Ok(Value::null());
|
||||
}
|
||||
|
||||
let mut sorted = state;
|
||||
|
|
@ -28,11 +29,11 @@ impl AggFunc for Median {
|
|||
|
||||
let len = sorted.len();
|
||||
if len % 2 == 1 {
|
||||
Value::from_float(sorted[len / 2])
|
||||
Ok(Value::from_float(sorted[len / 2]))
|
||||
} else {
|
||||
let mid1 = sorted[len / 2 - 1];
|
||||
let mid2 = sorted[len / 2];
|
||||
Value::from_float((mid1 + mid2) / 2.0)
|
||||
Ok(Value::from_float((mid1 + mid2) / 2.0))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -41,8 +42,8 @@ impl AggFunc for Median {
|
|||
struct Percentile;
|
||||
|
||||
impl AggFunc for Percentile {
|
||||
type State = (Vec<f64>, Option<f64>, Option<&'static str>);
|
||||
|
||||
type State = (Vec<f64>, Option<f64>, Option<Self::Error>);
|
||||
type Error = &'static str;
|
||||
const NAME: &'static str = "percentile";
|
||||
const ARGS: i32 = 2;
|
||||
|
||||
|
|
@ -69,16 +70,16 @@ impl AggFunc for Percentile {
|
|||
}
|
||||
}
|
||||
|
||||
fn finalize(state: Self::State) -> Value {
|
||||
fn finalize(state: Self::State) -> Result<Value, Self::Error> {
|
||||
let (mut values, p_value, err_value) = state;
|
||||
if values.is_empty() {
|
||||
return Value::null();
|
||||
return Ok(Value::null());
|
||||
}
|
||||
if let Some(err) = err_value {
|
||||
return Value::error_with_message(err.into());
|
||||
return Err(err);
|
||||
}
|
||||
if values.len() == 1 {
|
||||
return Value::from_float(values[0]);
|
||||
return Ok(Value::from_float(values[0]));
|
||||
}
|
||||
|
||||
let p = p_value.unwrap();
|
||||
|
|
@ -89,10 +90,12 @@ impl AggFunc for Percentile {
|
|||
let upper = index.ceil() as usize;
|
||||
|
||||
if lower == upper {
|
||||
Value::from_float(values[lower])
|
||||
Ok(Value::from_float(values[lower]))
|
||||
} else {
|
||||
let weight = index - lower as f64;
|
||||
Value::from_float(values[lower] * (1.0 - weight) + values[upper] * weight)
|
||||
Ok(Value::from_float(
|
||||
values[lower] * (1.0 - weight) + values[upper] * weight,
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -101,8 +104,8 @@ impl AggFunc for Percentile {
|
|||
struct PercentileCont;
|
||||
|
||||
impl AggFunc for PercentileCont {
|
||||
type State = (Vec<f64>, Option<f64>, Option<&'static str>);
|
||||
|
||||
type State = (Vec<f64>, Option<f64>, Option<Self::Error>);
|
||||
type Error = &'static str;
|
||||
const NAME: &'static str = "percentile_cont";
|
||||
const ARGS: i32 = 2;
|
||||
|
||||
|
|
@ -129,16 +132,16 @@ impl AggFunc for PercentileCont {
|
|||
}
|
||||
}
|
||||
|
||||
fn finalize(state: Self::State) -> Value {
|
||||
fn finalize(state: Self::State) -> Result<Value, Self::Error> {
|
||||
let (mut values, p_value, err_state) = state;
|
||||
if values.is_empty() {
|
||||
return Value::null();
|
||||
return Ok(Value::null());
|
||||
}
|
||||
if let Some(err) = err_state {
|
||||
return Value::error_with_message(err.into());
|
||||
return Err(err);
|
||||
}
|
||||
if values.len() == 1 {
|
||||
return Value::from_float(values[0]);
|
||||
return Ok(Value::from_float(values[0]));
|
||||
}
|
||||
|
||||
let p = p_value.unwrap();
|
||||
|
|
@ -149,10 +152,12 @@ impl AggFunc for PercentileCont {
|
|||
let upper = index.ceil() as usize;
|
||||
|
||||
if lower == upper {
|
||||
Value::from_float(values[lower])
|
||||
Ok(Value::from_float(values[lower]))
|
||||
} else {
|
||||
let weight = index - lower as f64;
|
||||
Value::from_float(values[lower] * (1.0 - weight) + values[upper] * weight)
|
||||
Ok(Value::from_float(
|
||||
values[lower] * (1.0 - weight) + values[upper] * weight,
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -161,8 +166,8 @@ impl AggFunc for PercentileCont {
|
|||
struct PercentileDisc;
|
||||
|
||||
impl AggFunc for PercentileDisc {
|
||||
type State = (Vec<f64>, Option<f64>, Option<&'static str>);
|
||||
|
||||
type State = (Vec<f64>, Option<f64>, Option<Self::Error>);
|
||||
type Error = &'static str;
|
||||
const NAME: &'static str = "percentile_disc";
|
||||
const ARGS: i32 = 2;
|
||||
|
||||
|
|
@ -170,19 +175,19 @@ impl AggFunc for PercentileDisc {
|
|||
Percentile::step(state, args);
|
||||
}
|
||||
|
||||
fn finalize(state: Self::State) -> Value {
|
||||
fn finalize(state: Self::State) -> Result<Value, Self::Error> {
|
||||
let (mut values, p_value, err_value) = state;
|
||||
if values.is_empty() {
|
||||
return Value::null();
|
||||
return Ok(Value::null());
|
||||
}
|
||||
if let Some(err) = err_value {
|
||||
return Value::error_with_message(err.into());
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
let p = p_value.unwrap();
|
||||
values.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
||||
let n = values.len() as f64;
|
||||
let index = (p * (n - 1.0)).floor() as usize;
|
||||
Value::from_float(values[index])
|
||||
Ok(Value::from_float(values[index]))
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
use limbo_ext::{
|
||||
register_extension, ExtensionApi, ResultCode, VTabCursor, VTabModule, VTabModuleDerive, Value,
|
||||
register_extension, ResultCode, VTabCursor, VTabKind, VTabModule, VTabModuleDerive, Value,
|
||||
};
|
||||
|
||||
register_extension! {
|
||||
|
|
@ -16,36 +16,38 @@ macro_rules! try_option {
|
|||
}
|
||||
|
||||
/// A virtual table that generates a sequence of integers
|
||||
#[derive(Debug, VTabModuleDerive)]
|
||||
#[derive(Debug, VTabModuleDerive, Default)]
|
||||
struct GenerateSeriesVTab;
|
||||
|
||||
impl VTabModule for GenerateSeriesVTab {
|
||||
type VCursor = GenerateSeriesCursor;
|
||||
type Error = ResultCode;
|
||||
const NAME: &'static str = "generate_series";
|
||||
const VTAB_KIND: VTabKind = VTabKind::TableValuedFunction;
|
||||
|
||||
fn connect(api: &ExtensionApi) -> ResultCode {
|
||||
fn create_schema(_args: &[Value]) -> String {
|
||||
// Create table schema
|
||||
let sql = "CREATE TABLE generate_series(
|
||||
"CREATE TABLE generate_series(
|
||||
value INTEGER,
|
||||
start INTEGER HIDDEN,
|
||||
stop INTEGER HIDDEN,
|
||||
step INTEGER HIDDEN
|
||||
)";
|
||||
api.declare_virtual_table(Self::NAME, sql)
|
||||
)"
|
||||
.into()
|
||||
}
|
||||
|
||||
fn open() -> Self::VCursor {
|
||||
GenerateSeriesCursor {
|
||||
fn open(&self) -> Result<Self::VCursor, Self::Error> {
|
||||
Ok(GenerateSeriesCursor {
|
||||
start: 0,
|
||||
stop: 0,
|
||||
step: 0,
|
||||
current: 0,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn filter(cursor: &mut Self::VCursor, arg_count: i32, args: &[Value]) -> ResultCode {
|
||||
fn filter(cursor: &mut Self::VCursor, args: &[Value]) -> ResultCode {
|
||||
// args are the start, stop, and step
|
||||
if arg_count == 0 || arg_count > 3 {
|
||||
if args.is_empty() || args.len() > 3 {
|
||||
return ResultCode::InvalidArgs;
|
||||
}
|
||||
let start = try_option!(args[0].to_integer(), ResultCode::InvalidArgs);
|
||||
|
|
@ -78,7 +80,7 @@ impl VTabModule for GenerateSeriesVTab {
|
|||
ResultCode::OK
|
||||
}
|
||||
|
||||
fn column(cursor: &Self::VCursor, idx: u32) -> Value {
|
||||
fn column(cursor: &Self::VCursor, idx: u32) -> Result<Value, Self::Error> {
|
||||
cursor.column(idx)
|
||||
}
|
||||
|
||||
|
|
@ -163,14 +165,14 @@ impl VTabCursor for GenerateSeriesCursor {
|
|||
false
|
||||
}
|
||||
|
||||
fn column(&self, idx: u32) -> Value {
|
||||
match idx {
|
||||
fn column(&self, idx: u32) -> Result<Value, Self::Error> {
|
||||
Ok(match idx {
|
||||
0 => Value::from_integer(self.current),
|
||||
1 => Value::from_integer(self.start),
|
||||
2 => Value::from_integer(self.stop),
|
||||
3 => Value::from_integer(self.step),
|
||||
_ => Value::null(),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn rowid(&self) -> i64 {
|
||||
|
|
@ -227,7 +229,8 @@ mod tests {
|
|||
}
|
||||
// Helper function to collect all values from a cursor, returns Result with error code
|
||||
fn collect_series(series: Series) -> Result<Vec<i64>, ResultCode> {
|
||||
let mut cursor = GenerateSeriesVTab::open();
|
||||
let tbl = GenerateSeriesVTab;
|
||||
let mut cursor = tbl.open()?;
|
||||
|
||||
// Create args array for filter
|
||||
let args = vec![
|
||||
|
|
@ -237,7 +240,7 @@ mod tests {
|
|||
];
|
||||
|
||||
// Initialize cursor through filter
|
||||
match GenerateSeriesVTab::filter(&mut cursor, 3, &args) {
|
||||
match GenerateSeriesVTab::filter(&mut cursor, &args) {
|
||||
ResultCode::OK => (),
|
||||
ResultCode::EOF => return Ok(vec![]),
|
||||
err => return Err(err),
|
||||
|
|
@ -245,7 +248,7 @@ mod tests {
|
|||
|
||||
let mut values = Vec::new();
|
||||
loop {
|
||||
values.push(cursor.column(0).to_integer().unwrap());
|
||||
values.push(cursor.column(0)?.to_integer().unwrap());
|
||||
if values.len() > 1000 {
|
||||
panic!(
|
||||
"Generated more than 1000 values, expected this many: {:?}",
|
||||
|
|
@ -543,8 +546,8 @@ mod tests {
|
|||
let start = series.start;
|
||||
let stop = series.stop;
|
||||
let step = series.step;
|
||||
|
||||
let mut cursor = GenerateSeriesVTab::open();
|
||||
let tbl = GenerateSeriesVTab::default();
|
||||
let mut cursor = tbl.open().unwrap();
|
||||
|
||||
let args = vec![
|
||||
Value::from_integer(start),
|
||||
|
|
@ -553,7 +556,7 @@ mod tests {
|
|||
];
|
||||
|
||||
// Initialize cursor through filter
|
||||
GenerateSeriesVTab::filter(&mut cursor, 3, &args);
|
||||
GenerateSeriesVTab::filter(&mut cursor, &args);
|
||||
|
||||
let mut rowids = vec![];
|
||||
while !GenerateSeriesVTab::eof(&cursor) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue