diff --git a/src/algorithm/embeddings/fastrp.rs b/src/algorithm/embeddings/fastrp.rs index 56d209e..757f769 100644 --- a/src/algorithm/embeddings/fastrp.rs +++ b/src/algorithm/embeddings/fastrp.rs @@ -419,7 +419,13 @@ impl<'a> FastRPBuilder<'a> { // remaining iterates fold in one at a time: acc = acc + w_t · H_t. // The full outer join grows the key set as iterates reach new // vertices (an absent iterate row is a zero contribution). - for (t, w) in kept.iter().skip(2) { + // + // Every step gets a *unique* checkpoint name: `push_pre_sorted` + // registers the directory in the eviction queue, and a repeated + // name would enqueue the same directory twice — the next + // `evict_all_but_latest_n` would then delete the directory the + // freshly returned (lazy) accumulator still scans. + for (step, (t, w)) in kept.iter().skip(2).enumerate() { let side = states[*t].clone().select(vec![ col(VERTEX_ID).alias("__fastrp_ht_vid"), col(EMBEDDING).alias("__fastrp_ht"), @@ -440,7 +446,7 @@ impl<'a> FastRPBuilder<'a> { .alias(ACC), ])?; acc = folds_checkpointer - .push_pre_sorted(&ctx, "fold", acc.clone(), ACC_VID) + .push_pre_sorted(&ctx, &format!("fold-{}", step + 1), acc.clone(), ACC_VID) .await?; // only the latest accumulator checkpoint is needed folds_checkpointer.evict_all_but_latest_n(&ctx, 1).await?; @@ -482,6 +488,7 @@ impl<'a> FastRPBuilder<'a> { edges_checkpointer.purge(&ctx).await?; states_checkpointer.purge(&ctx).await?; vertices_checkpointer.purge(&ctx).await?; + folds_checkpointer.purge(&ctx).await?; log::info!( "FastRP {run_id} finished after {} iterations, output: {output}", @@ -803,6 +810,65 @@ mod tests { Ok(()) } + /// Regression test: the default `iterations(4)` + all-ones weights fold + /// K = 4 iterates, so the checkpointed fold loop runs twice. The loop used + /// to reuse the checkpoint name `fold` for every step; the second + /// `evict_all_but_latest_n` then deleted the directory the final join was + /// about to scan (`NotFound .../folds/fold/part-*.parquet`). + #[tokio::test] + async fn fastrp_four_iterations_leave_no_stale_fold_checkpoints() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("fold4"); + let d = 2; + let seed = 42u64; + // cycle 1 -> 2 -> 3 -> 1: H_t(v) = init(pred^t(v)) + let g = create_graph(vec![1, 2, 3], vec![vec![1, 2], vec![2, 3], vec![3, 1]])?; + + g.fastrp() // default iteration_weights = all ones + .dim(d) + .iterations(4) + .seed(seed) + .set_checkpoint_dir(checkpoint_dir.clone()) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + let i1 = init_ref(1, seed, d); + let i2 = init_ref(2, seed, d); + let i3 = init_ref(3, seed, d); + + // H(1) = i3 + i2 + i1 + i3, H(2) = i1 + i3 + i2 + i1, H(3) = i2 + i1 + i3 + i2 + assert_eq!(got[&1], add(&add(&add(&i3, &i2), &i1), &i3)); + assert_eq!(got[&2], add(&add(&add(&i1, &i3), &i2), &i1)); + assert_eq!(got[&3], add(&add(&add(&i2, &i1), &i3), &i2)); + + // the run must not leak checkpoint files (empty dirs may remain) + fn count_files(dir: &std::path::Path) -> usize { + let mut n = 0; + for entry in fs::read_dir(dir).into_iter().flatten().flatten() { + if entry.file_type().map(|t| t.is_dir()).unwrap_or(false) { + n += count_files(&entry.path()); + } else { + n += 1; + } + } + n + } + let checkpoint_root: std::path::PathBuf = checkpoint_dir + .to_string() + .trim_start_matches("file://") + .into(); + assert_eq!( + count_files(&checkpoint_root), + 0, + "checkpoint files must be purged after a successful run" + ); + Ok(()) + } + #[tokio::test] async fn fastrp_weights_length_is_validated() -> Result<()> { let (ctx, checkpoint_dir, output_uri, _guard) = setup("weights_len"); diff --git a/src/ml/kmeans.rs b/src/ml/kmeans.rs index a030b17..067a3a2 100644 --- a/src/ml/kmeans.rs +++ b/src/ml/kmeans.rs @@ -780,7 +780,6 @@ mod tests { use datafusion::arrow::datatypes::{Field, Schema}; use datafusion::datasource::MemTable; use std::sync::Arc; - use std::time::Instant; fn lcg(seed: &mut u32) -> f32 { *seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);