Add new function to do bool operation on dict_id header

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) {

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<Option<u32>, 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<u8> {
+        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);
+    }
+}

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(())
+}

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;

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(())
 }

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]