Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 17 additions & 16 deletions sqlite-compress/src/dict/train.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,23 +106,16 @@ pub fn train_by_column<C: SetupConnection, F: FnMut(TrainProgress)>(

let state = ExtensionState::from_db(unsafe { connection.sqlite_handle() })?;

let (enough, available) =
has_enough_samples(&state, &table_name, &column_name, column.min_samples)?;
let available = has_enough_samples(&state, &table_name, &column_name, column.min_samples)?;

if !enough {
return Ok(None);
}

if !has_retrain_required(
has_retrain_required(
&state,
&dict_store,
key.table(),
key.column(),
available,
column.retrain_growth,
)? {
return Ok(None);
}
)?;

let dictionary = build_dictionary(
&state,
Expand Down Expand Up @@ -194,7 +187,7 @@ fn has_enough_samples<C: AsRef<Connection>>(
table_name: &str,
column_name: &str,
min_samples: usize,
) -> Result<(bool, i64), SetupError> {
) -> Result<i64, SetupError> {
let connection = connection.as_ref();

let rows = connection
Expand All @@ -209,7 +202,13 @@ fn has_enough_samples<C: AsRef<Connection>>(

let count = first_value(&rows).and_then(SqlValue::as_i64).unwrap_or(0);

Ok((count >= min_samples as i64, count))
if count >= min_samples as i64 {
return Ok(count);
}

Err(SetupError::InvalidConfig(
"not enough samples available: {count} < {min_samples}",
))
}

/// Check if the retrain is required.
Expand All @@ -224,7 +223,7 @@ fn has_retrain_required<C: AsRef<Connection>>(
column_name: &str,
available: i64,
retrain_growth: usize,
) -> Result<bool, SetupError> {
) -> Result<(), SetupError> {
let connection = connection.as_ref();

let rows = connection
Expand All @@ -245,11 +244,13 @@ fn has_retrain_required<C: AsRef<Connection>>(

let last = first_value(&rows).and_then(SqlValue::as_i64).unwrap_or(0);

if last < 0 {
return Ok(true);
if last < 0 || available >= last + retrain_growth as i64 {
return Ok(());
}

Ok(available >= last + retrain_growth as i64)
Err(SetupError::InvalidConfig(
"not enough samples available: {available} < {last} + {retrain_growth}",
))
}

/// Samples fetched per random query. Scales with `max_samples` so we keep
Expand Down
18 changes: 13 additions & 5 deletions sqlite-compress/tests/dict_train.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use sqlite_compress::{
DEFAULT_LEVEL, DEFAULT_RETRAIN_GROWTH, ExtensionState, Header, SchemaName, SetupColumn,
SetupConfig, SetupTable, TableName, TrainOptions, TrainProgress, setup, train_all,
SetupConfig, SetupError, SetupTable, TableName, TrainOptions, TrainProgress, setup, train_all,
train_by_column,
};

Expand Down Expand Up @@ -247,8 +247,11 @@ fn train_skips_column_without_enough_samples() {

setup(&db, &config).unwrap();

let dict_ids = train_all(&db, &config, 1024, TrainOptions::default()).unwrap();
assert_eq!(dict_ids.len(), 1);
let err = train_all(&db, &config, 1024, TrainOptions::default()).unwrap_err();
assert!(matches!(
err,
SetupError::InvalidConfig("not enough samples available: {count} < {min_samples}")
));

let rows = state
.as_ref()
Expand Down Expand Up @@ -460,8 +463,13 @@ fn train_skips_when_growth_below_threshold() {
// +100 samples << RETRAIN_GROWTH (5000)
insert_samples(&state, 100);

let second = train_all(&db, &config, 1024, TrainOptions::default()).unwrap();
assert!(second.is_empty());
let err = train_all(&db, &config, 1024, TrainOptions::default()).unwrap_err();
assert!(matches!(
err,
SetupError::InvalidConfig(
"not enough samples available: {available} < {last} + {retrain_growth}"
)
));

let rows = state
.as_ref()
Expand Down
Loading