crates/mimir-ingest/tests/ingest/compact.rs
//! The lake's compaction job (`EventLog::compact`): every uncompressed or
//! small data file of a table rewritten as zstd and merged, in ONE `replace`
//! snapshot per table — the same rows before and after, the cache stamp
//! unmoved, and a crash between writing the new files and committing them
//! leaves a lake every reader still reads whole.
use std::path::{Path, PathBuf};
use mimir_ingest::{CompactOptions, Consumer, EventLog, Kind, RepoSpec};
use skade::arrow_array::RecordBatch;
use skade::parquet::basic::Compression;
use skade::parquet::file::reader::{FileReader, SerializedFileReader};
use crate::common::{FakeDumper, Scratch, open_log};
const REPO: &str = "r";
fn build_history(src: &Scratch) -> Vec<String> {
let mut shas = Vec::new();
src.write("crates/a/src/lib.rs", "fn f (x: i32) | one\nfn g () | two\ncall g -> f\n")
.write("crates/b/src/lib.rs", "fn h () | three\ncall h -> h\n");
shas.push(src.commit("one"));
src.write("crates/a/src/lib.rs", "fn f (x: i32) | uno\nfn g () | two\ncall g -> f\ncall f -> g\n");
shas.push(src.commit("two"));
src.write("crates/a/src/lib.rs", "fn f (x: i64) | uno\nfn g () | two\ncall f -> g\n");
src.remove("crates/b");
shas.push(src.commit("three"));
shas
}
fn spec(tmp: &Path, src: &Scratch) -> RepoSpec {
RepoSpec {
name: REPO.into(),
source: src.root.to_string_lossy().into_owned(),
branch: "main".into(),
worktree: tmp.join("clone"),
initial_offset: None,
}
}
/// Every row of every event table, as sorted canonical byte rows (order-free).
async fn all_rows(log: &EventLog) -> Vec<Vec<u8>> {
let mut out = Vec::new();
for kind in Kind::ALL {
for b in log.events(kind).await.unwrap() {
out.extend(rows_of(&b));
}
}
out.sort();
out
}
fn rows_of(b: &RecordBatch) -> Vec<Vec<u8>> {
use arrow_row::{RowConverter, SortField};
let fields: Vec<SortField> = b.schema().fields().iter().map(|f| SortField::new(f.data_type().clone())).collect();
let conv = RowConverter::new(fields).unwrap();
let rows = conv.convert_columns(b.columns()).unwrap();
(0..b.num_rows()).map(|i| rows.row(i).as_ref().to_vec()).collect()
}
fn data_files(lake: &Path, table: &str) -> Vec<PathBuf> {
let dir = lake.join("warehouse/mimir_v2").join(table).join("data");
let mut v: Vec<PathBuf> = match std::fs::read_dir(&dir) {
Ok(rd) => rd.map(|e| e.unwrap().path()).filter(|p| p.extension().is_some_and(|x| x == "parquet")).collect(),
Err(_) => Vec::new(),
};
v.sort();
v
}
fn codecs(files: &[PathBuf]) -> Vec<Compression> {
files
.iter()
.map(|p| {
let r = SerializedFileReader::new(std::fs::File::open(p).unwrap()).unwrap();
r.metadata().row_group(0).column(0).compression()
})
.collect()
}
fn is_zstd(c: &Compression) -> bool {
matches!(c, Compression::ZSTD(_))
}
async fn consumed_lake(tmp: &Path) -> (EventLog, Vec<String>, Scratch) {
let src = Scratch::new(tmp.join("src"));
let shas = build_history(&src);
let log = open_log(&tmp.join("lake")).await;
let mut c = Consumer::boot(log, FakeDumper::default()).await.unwrap();
c.run(&spec(tmp, &src)).await.unwrap();
let (log, _) = c.into_parts();
(log, shas, src)
}
#[tokio::test(flavor = "current_thread")]
async fn compaction_keeps_every_row_and_moves_no_stamp() {
let tmp = tempfile::tempdir().unwrap();
let lake = tmp.path().join("lake");
let (log, _shas, src) = consumed_lake(tmp.path()).await;
// three releases → three small data files per event table that got events
// (the consumer writes zstd; the lake's old uncompressed files are the
// `select` unit test's and the live copy's case)
let before_files = data_files(&lake, "code_edge_events");
assert!(before_files.len() >= 2, "{before_files:?}");
assert!(codecs(&before_files).iter().all(is_zstd), "{:?}", codecs(&before_files));
let rows_before = all_rows(&log).await;
let fold_before = log.fold(None, None).await.unwrap();
let stamp_before = log.stamp().await.unwrap();
let runs_before = log.runs(REPO).await.unwrap();
let opts = CompactOptions { grace: std::time::Duration::ZERO, keep_snapshots: 1, ..Default::default() };
let reports = log.compact(&opts).await.unwrap();
let edges = reports.iter().find(|r| r.table == "code_edge_events").unwrap();
assert!(edges.files_in >= 2 && edges.files_out == 1, "{edges}");
assert_eq!(edges.rows, rows_of_kind(&log, Kind::Edges).await, "{edges}");
assert!(edges.snapshot.is_some(), "a replace snapshot was committed: {edges}");
// the same rows, the same fold, the same ledger — in zstd files, merged
assert_eq!(all_rows(&log).await, rows_before);
assert_eq!(log.fold(None, None).await.unwrap(), fold_before);
assert_eq!(log.runs(REPO).await.unwrap(), runs_before);
for table in ["code_node_events", "code_edge_events", "code_span_events", "ingest_runs", "consumer_offsets"] {
let files = data_files(&lake, table);
assert_eq!(files.len(), 1, "{table}: {files:?}");
assert!(codecs(&files).iter().all(is_zstd), "{table}: {:?}", codecs(&files));
}
// the cache stamp did not move: a replace adds no event
assert_eq!(log.stamp().await.unwrap(), stamp_before);
// only the replace is kept (the stamp above still names the append it replaced)
let t = log.warehouse().table("mimir_v2.code_edge_events").await.unwrap();
assert_eq!(t.inner().metadata().snapshots().count(), 1);
// and its summary counts the rows once: a replace removes what it adds
let summary = t.inner().metadata().current_snapshot().unwrap().summary().additional_properties.clone();
assert_eq!(summary.get("total-records").map(String::as_str), Some(edges.rows.to_string().as_str()), "{summary:?}");
assert_eq!(summary.get("total-data-files").map(String::as_str), Some("1"), "{summary:?}");
// a second run finds nothing to do
let again = log.compact(&opts).await.unwrap();
assert!(again.iter().all(|r| r.files_out == 0 && r.snapshot.is_none()), "{again:?}");
// the log goes on: a new release appends after the replace, and folds
src.write("crates/a/src/lib.rs", "fn f (x: i64) | dos\nfn g () | two\ncall f -> g\n");
src.commit("four");
let mut c = Consumer::boot(log, FakeDumper::default()).await.unwrap();
let reports = c.run(&spec(tmp.path(), &src)).await.unwrap();
assert_eq!(reports.len(), 1);
let (log, _) = c.into_parts();
assert_ne!(log.stamp().await.unwrap(), stamp_before);
let at4 = log.fold(Some(REPO), None).await.unwrap();
assert_ne!(at4.repo(Kind::Nodes, REPO).unwrap().get("a:lib", "a f"), fold_before.repo(Kind::Nodes, REPO).unwrap().get("a:lib", "a f"));
}
async fn rows_of_kind(log: &EventLog, kind: Kind) -> u64 {
log.events(kind).await.unwrap().iter().map(|b| b.num_rows() as u64).sum()
}
#[tokio::test(flavor = "current_thread")]
async fn a_crash_after_writing_the_new_files_leaves_a_readable_lake() {
let tmp = tempfile::tempdir().unwrap();
let lake = tmp.path().join("lake");
let (log, _shas, _src) = consumed_lake(tmp.path()).await;
let rows_before = all_rows(&log).await;
let stamp_before = log.stamp().await.unwrap();
// the first table the job rewrites is where it dies
let live_before = data_files(&lake, "code_node_events");
// the job dies after the new files are in storage, before the commit
let crash = CompactOptions { grace: std::time::Duration::ZERO, fail_before_commit: true, ..Default::default() };
let err = log.compact(&crash).await.unwrap_err();
assert!(format!("{err:#}").contains("fail_before_commit"), "{err:#}");
// nothing moved for a reader: same snapshot, same rows; the new files are orphans
assert_eq!(log.stamp().await.unwrap(), stamp_before);
assert_eq!(all_rows(&log).await, rows_before);
let after_crash = data_files(&lake, "code_node_events");
assert!(after_crash.len() > live_before.len(), "orphans left: {after_crash:?}");
// the next run completes and sweeps the orphans
let opts = CompactOptions { grace: std::time::Duration::ZERO, keep_snapshots: 1, ..Default::default() };
let reports = log.compact(&opts).await.unwrap();
let nodes = reports.iter().find(|r| r.table == "code_node_events").unwrap();
assert!(nodes.swept_files >= after_crash.len(), "the orphans and the replaced files: {nodes}");
let files = data_files(&lake, "code_node_events");
assert_eq!(files.len(), 1, "{files:?}");
assert!(codecs(&files).iter().all(is_zstd));
assert_eq!(all_rows(&log).await, rows_before);
assert_eq!(log.stamp().await.unwrap(), stamp_before);
}