diff --git a/crates/sqlite-ffi/src/context/mod.rs b/crates/sqlite-ffi/src/context/mod.rs index 7730d37..7c68c47 100644 --- a/crates/sqlite-ffi/src/context/mod.rs +++ b/crates/sqlite-ffi/src/context/mod.rs @@ -4,6 +4,7 @@ use std::{ffi::c_void, marker::PhantomData, ptr::NonNull}; use libsqlite3_sys::{ sqlite3_context, sqlite3_result_blob, sqlite3_result_error, sqlite3_result_int64, + sqlite3_result_null, }; pub use crate::context::aux::Aux; @@ -52,6 +53,12 @@ impl Context { pub fn result_int64(&self, value: i64) { unsafe { result_int_raw(self.ctx.as_ptr(), value) }; } + + pub fn result_null(&self) { + unsafe { + sqlite3_result_null(self.ctx.as_ptr()); + } + } } unsafe fn result_blob_raw(context: *mut sqlite3_context, blob: *const c_void, size: i32) { diff --git a/sqlite-compress/src/functions/compressed_dict_id.rs b/sqlite-compress/src/functions/compressed_dict_id.rs new file mode 100644 index 0000000..60f6d42 --- /dev/null +++ b/sqlite-compress/src/functions/compressed_dict_id.rs @@ -0,0 +1,60 @@ +use sqlite_ffi::{Context, SqliteError, Value}; + +use crate::functions::{ParsedBlob, errors::CodecError, header::Header}; + +/// SQLite callback for `compressed_dict_id(blob)`. +/// +/// Returns the dictionary ID stored in the blob header, or NULL if the blob +/// is not in a recognized compressed format. +pub fn sqlite_compressed_dict_id(context: Context, values: &[Value]) -> Result<(), SqliteError> { + match dict_id_from_blob(values[0].to_blob()).map_err(|e| SqliteError::Message(e.to_string()))? { + Some(dict_id) => context.result_int64(i64::from(dict_id)), + None => context.result_null(), + } + + Ok(()) +} + +fn dict_id_from_blob(blob: &[u8]) -> Result, CodecError> { + match Header::parse_blob(blob)? { + ParsedBlob::Framed(header, _) => Ok(Some(header.dict_id.get())), + ParsedBlob::Plain => Ok(None), + } +} + +#[cfg(test)] +mod tests { + use super::dict_id_from_blob; + use crate::{DictId, functions::Header}; + + fn framed_blob(dict_id: u32) -> Vec { + let header = Header::new(DictId::from(dict_id), 0); + let mut blob = Vec::new(); + header.write(&mut blob); + blob + } + + #[test] + fn returns_dictionary_id_from_framed_blob() { + let blob = framed_blob(42); + + assert_eq!(dict_id_from_blob(&blob).unwrap(), Some(42)); + } + + #[test] + fn returns_zero_for_raw_zstd_blob() { + let blob = framed_blob(0); + + assert_eq!(dict_id_from_blob(&blob).unwrap(), Some(0)); + } + + #[test] + fn returns_none_for_plain_blob() { + assert_eq!(dict_id_from_blob(b"plain text").unwrap(), None); + } + + #[test] + fn returns_none_for_malformed_header() { + assert_eq!(dict_id_from_blob(&[0, 1, 2]).unwrap(), None); + } +} diff --git a/sqlite-compress/src/functions/latest_dict_id.rs b/sqlite-compress/src/functions/latest_dict_id.rs new file mode 100644 index 0000000..99cdc2f --- /dev/null +++ b/sqlite-compress/src/functions/latest_dict_id.rs @@ -0,0 +1,22 @@ +use sqlite_ffi::{Context, SqliteError, Value}; + +use crate::{ColumnKey, dict::read_current_id, state::ExtensionState}; + +pub fn sqlite_latest_dict_id(context: Context, values: &[Value]) -> Result<(), SqliteError> { + let schema = values[0].to_text()?; + let table = values[1].to_text()?; + let column = values[2].to_text()?; + + let state = ExtensionState::from_context(&context)?; + let key = ColumnKey::new(schema, table, column); + + let latest_id = + read_current_id(&state, &key).map_err(|error| SqliteError::Message(error.to_string()))?; + + match latest_id { + Some(dict_id) => context.result_int64(i64::from(dict_id.get())), + None => context.result_null(), + } + + Ok(()) +} diff --git a/sqlite-compress/src/functions/mod.rs b/sqlite-compress/src/functions/mod.rs index 68689da..c1e5c2d 100644 --- a/sqlite-compress/src/functions/mod.rs +++ b/sqlite-compress/src/functions/mod.rs @@ -1,7 +1,9 @@ mod compress; +mod compressed_dict_id; mod decompress; mod errors; mod header; +mod latest_dict_id; #[cfg(feature = "like")] mod like; mod types; @@ -10,11 +12,13 @@ mod types; pub const DEFAULT_LEVEL: types::Level = types::Level::new(3); pub use compress::sqlite_compress; +pub use compressed_dict_id::sqlite_compressed_dict_id; pub use decompress::sqlite_decompress; pub(crate) use decompress::{decompress_raw, decompress_with_decoder}; pub use errors::CodecError; pub use header::Header; pub(crate) use header::ParsedBlob; +pub use latest_dict_id::sqlite_latest_dict_id; #[cfg(feature = "like")] pub use like::sqlite_decompress_like; pub use types::Level; diff --git a/sqlite-compress/src/lib.rs b/sqlite-compress/src/lib.rs index a71648b..eebd52b 100644 --- a/sqlite-compress/src/lib.rs +++ b/sqlite-compress/src/lib.rs @@ -37,5 +37,13 @@ pub fn sqlite3_compress_init(connection: Connection) -> Result<(), SqliteError> #[cfg(feature = "like")] connection.create_function("decompress_like", 6, functions::sqlite_decompress_like)?; + connection.create_function( + "compressed_dict_id", + 1, + functions::sqlite_compressed_dict_id, + )?; + + connection.create_function("latest_dict_id", 3, functions::sqlite_latest_dict_id)?; + Ok(()) } diff --git a/sqlite-compress/tests/dict_train.rs b/sqlite-compress/tests/dict_train.rs index c808c95..f05c3e5 100644 --- a/sqlite-compress/tests/dict_train.rs +++ b/sqlite-compress/tests/dict_train.rs @@ -812,7 +812,7 @@ fn second_connection_compress_uses_trained_dict_id() { fn train_by_column_reports_progress() { let sample_count = 64; let max_samples = 64; - let progress_every = 10; + let progress_every = 1; let db = common::TestDb::open(); let state = db.state(); @@ -867,7 +867,7 @@ fn train_by_column_reports_progress() { }), }; - let dict_id = train_by_column( + train_by_column( &db, &table, &column, @@ -877,15 +877,14 @@ fn train_by_column_reports_progress() { ) .unwrap(); - assert!(dict_id.is_some()); assert!(!reports.is_empty()); - assert!(reports.iter().all(|&(_, target)| target == max_samples)); assert!( reports .iter() - .all(|&(done, _)| done > 0 && done.is_multiple_of(progress_every)) + .all(|&(done, target)| done > 0 && done <= target && target == max_samples) ); assert!(reports.windows(2).all(|w| w[0].0 <= w[1].0)); + assert_eq!(reports.last(), Some(&(max_samples, max_samples))); } #[test]